# REINFORCE

In the previous notebook, we looked at [Deep Q Network](./dqn.ipynb) and [Double Deep Q Networks](./ddqn.ipynb) as methods to estimate the $Q$-function by leveraging the Bellman equation. In the present notebook, we meet the [REINFORCE algorithm](https://link.springer.com/article/10.1007/BF00992696), which makes it possible to learn optimal policies directly by performing gradient descent on the policy parameters.

To this end, we build on the code samples from the amazing [pytorch-examples repository](https://github.com/pytorch/examples/blob/master/reinforcement_learning/reinforce.py) and follow closely the highly instructive conceptual derivations from [OpenAI spinning up](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html#).

## Basic Policy Gradient Theorem

In the **Policy Gradient** approach, our goal is to find a parametric policy $\pi_\theta$ maximizing the expected total rewards $J(\pi_\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)]$ under the policy $\pi_\theta$, where $R(\tau)$ denotes the total reward along a trajectory $\tau$ sampled according to the policy $\pi_\theta$.

The basic idea of the **REINFORCE algorithm** or policy gradient methods is to update the policy parameters $\pi_\theta$ via gradient descent. To that end, we need to compute the gradient $\nabla_\theta J(\pi_\theta)$ of the expected total rewards $J(\pi_\theta)$ with respect to $\theta$. That is, we put 
$$\theta_{k + 1} = \theta_k + \alpha \nabla_\theta J(\pi_\theta)$$
for a suitable learning rate $\alpha > 0$.

To compute the gradient, $\nabla_\theta J(\pi_\theta)$ we let $p_\theta(\tau)$ denote the probability density of sampling a trajectory $\tau$ when following the parametric policy $\pi_\theta$.  Then, we reexpress $\nabla_\theta J(\pi_\theta)$ via the **log-derivative trick**
\begin{align}
\nabla_\theta J(\pi_\theta) &= \nabla_\theta \int_\tau R(\tau)p_\theta(\tau) \mathrm d \tau 
= \int_\tau \nabla_\theta R(\tau)p_\theta(\tau) \mathrm d \tau 
=\int_\tau R(\tau)\nabla_\theta\log (p_\theta(\tau))  p_\theta(\tau) \mathrm d \tau 
= \mathbb E_{\tau \sim \pi_\theta}\big[\nabla_\theta R(\tau)\log (p_\theta(\tau)) \big]
\end{align}

Next, we compute the probability of observing a trajectory $\tau = (s_0, s_1, \dots, s_T)$ of length $T \ge 1$ when starting at a state $s_0$ and following the policy $\pi_\theta$. This probability factorizes as 
$$p_\theta(\tau) = \prod_{t < T} p(s_{t + 1}\,|\,s_t, a_t) \pi_\theta(a_t|s_t),$$
where $p(s_{t + 1}\,|\, s_t, a_t)$ denotes the transition probability of the environment transitions from state $s_t$ to state $s_{t + 1}$ when performing the action $a_t$. Since this probability does not depend on the parameter $\theta$, we conclude that 
$$\nabla_\theta \log (p_\theta(\tau)) = \sum_{t < T} \nabla_\theta \log \pi_\theta(a_t|s_t).$$
Hence, we obtain the **Policy-Gradient Theorem**
\begin{align*}
\nabla_\theta J(\pi_\theta) &= \mathbb E_{\tau  \sim \pi_\theta}\Big[ R(\tau) \sum_{t < T} \nabla_\theta \log \pi_\theta(a_t|s_t)\Big].
\end{align*}

Loosely speaking, we increase the probability of an action $a_t$ in state $s_t$ by taking a gradient step with respect to the log-probability scaled with the total reward.

To summarize, the log-derivative trick made it possible to render the problem of finding the parameters for the optimal policy amenable the powerful gradient descent technique known from supervised learning. However, as noted in the [OpenAI tutorial](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html), there are two fundamental differences to the supervised setting.

1. In stochastic gradient descent, we compute gradients on samples taken from a fixed data distribution. In contrast, in the REINFORCE algorithm, we sample from the policy $\pi_\theta$, so that the sampling distribution changes as we update $\theta$.

2. It is tempting to view $-\mathbb E_{\pi_\theta}\Big[ R(\tau) \sum_{t < T}  \log \pi_\theta(a_t|s_t)\Big]$ as a loss function and track how quickly it decreases over several iterations. However, since the expectation is taken with respect to the $\theta$-dependent distribution $\pi_\theta$, its gradient is not related to $\nabla_\theta J(\pi_\theta)$.

## REINFORCE Algorithm

If we replace the expectation $\mathbb E_{\pi_\theta}$ by an average over Monte Carlo samples from an agent following the policy $\pi_\theta$, then it is straightforward to implement the [REINFORCE Algorithm](https://link.springer.com/article/10.1007/BF00992696) for optimizing the policy $\pi_\theta$.

1. Draw a mini-batch of sample trajectories $\tau^1, \dots, \tau^N$ from $\pi_\theta$
2. Compute the rewards $R(\tau^1), \dots, R(\tau^N)$
3. Compute the gradient step
$$\theta' = \theta + \alpha \frac1N \sum_{i \le N} R(\tau^i)\sum_{t < T}\nabla_\theta \log(\pi_\theta(a_t^i|s_t^i)).$$
4. Replace $\theta \leftarrow \theta'$ and go back to 1.

## REINFORCE Algorithm for the Cartpole Example

Next, we illustrate how the REINFORCE Algorithm is applied to the ``CartPole`` example. First, we load the environment.


```python
import gym
import torch

env = gym.make('CartPole-v1')
```

Next, we define the policy network as a multilayer perceptron with one hidden layer. The policy network takes a sample from the state space as input and returns logits for the actions of moving to the left or to the right.


```python
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

nactions = env.action_space.n
state_dim = env.observation_space.shape[0]
nhidden = 128
p_dropout = .6

class Policy(nn.Module):
    """Neural network parametrizing the policy
    """
    def __init__(self):
        """Initialize the policy network
        """
        super(Policy, self).__init__()
        self.affine1 = nn.Linear(state_dim, nhidden)
        self.dropout = nn.Dropout(p = p_dropout)
        self.affine2 = nn.Linear(nhidden, nactions)

        #save log-probs for the application in the REINFORCE module later
        self.saved_log_probs = []
        self.rewards = []

    def forward(self, s):
        """Compute action from state
    
        # Arguments
            s: input state
        # Result
            suggested action
        """
        s = self.affine1(s)
        s = self.dropout(s)
        s = F.relu(s)
        action_scores = self.affine2(s)
        return F.softmax(action_scores, dim = 1)
    
policy = Policy()
```

To draw an action from the current policy at a state $s$, we send the state $s$ through the policy network and sample from the corresponding probabilities.


```python
from torch.distributions import Categorical

def select_action(s):
    """Select action according to policy
    
    # Arguments
        s: state
    # Result
        selected action
    """
    s = torch.from_numpy(s).float().unsqueeze(0)
    probs = policy(s)
    m = Categorical(probs)
    a = m.sample()

    policy.saved_log_probs.append(m.log_prob(a))
    return a.item()
```


Each time after finishing an environment, we apply the REINFORCE algorithm.


```python
def finish_episode():
    """Apply optimizer to policy network after each episode
    """
    R = 0
    policy_loss = []


    for r in policy.rewards[::-1]:
        R = r + GAMMA * R

    for log_prob in policy.saved_log_probs:
        policy_loss.append(-log_prob * R)

    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum()
    policy_loss.backward()
    optimizer.step()

    del policy.rewards[:]
    del policy.saved_log_probs[:]
```

When running the REINFORCE algorithm,  we see that  initially the rewards increase. However, we also notice substantial fluctuations in the course of the optimization.



```python
running_reward = 10
MAX_STEPS = int(1e3)
PRINT_FREQ = int(1e3)
MA_WEIGHT = .05

def reinforce(env,
             neps = int(1e4),
             init_reward = 10):
    """Apply reinforce algorithm for several episodes
    
    # Arguments
        env: environment
        neps: number of episodes
        init_reward: initialization of running reward
    """
    running_reward = init_reward
    
    for i in range(neps):
        s, ep_reward = env.reset(), 0

        for _ in range(MAX_STEPS):  
            #perform step in environment
            a = select_action(s)
            s, r, d, _ = env.step(a)
            
            #collect rewards
            policy.rewards.append(r)
            ep_reward += r
            if d:
                break

        running_reward = MA_WEIGHT * ep_reward + (1 - MA_WEIGHT) * running_reward
        if i % PRINT_FREQ == 0: print(running_reward)
        finish_episode()
```

## Reward-to-Go Policy Gradient

As mentioned above, in the policy-gradient theorem
\begin{align*}
\nabla_\theta J(\pi_\theta) &= \mathbb E_{\tau  \sim \pi_\theta}\Big[ R(\tau) \sum_{t < T} \nabla_\theta \log \pi_\theta(a_t|s_t)\Big],
\end{align*}
the gradient of the log-probability of an action is scaled by the total reward along a trajectory. However, this is counter-intuitive. Indeed, the probability of taking an action $a_t$ at time $t$ should not depend on rewards before time $t$.

To make this precise, we expand $R(\tau)$ as $R(\tau) = \sum_{t' < T} \gamma^{t'} r_{t'}$, where $\gamma \in (0,1)$ is the discount factor and $r_{t'} = r(s_{t'}, a_{t'})$ denotes the reward obtained by applying the action $a_{t'}$ in state $s_{t'}$. Hence,
\begin{align*}
\nabla_\theta J(\pi_\theta) &=  \sum_{t < T}  \sum_{t' < T} \mathbb E_{\tau \sim \pi_\theta}\Big[ \gamma^{t'}r_{t'}\nabla_\theta \log \pi_\theta(a_t|s_t)\Big].
\end{align*}
We claim that in this double-sum, the terms with $t' < t$ vanish.

Indeed, conditioning on $s_{t'}$ and $a_{t'}$ gives that 
\begin{align*}
\mathbb E_{\tau \sim \pi_\theta}\Big[ r_{t'}\nabla_\theta \log \pi_\theta(a_t|s_t)\Big] &=
\int r(s_{t'}, a_{t'}) \nabla_\theta \log \pi_\theta(a_t|s_t) p_\theta(s_{t'}, a_{t'}, s_t, a_t) \mathrm d \tau\\
&= \int r(s_{t'}, a_{t'}) \Big(\int \nabla_\theta \log(\pi_\theta(a_t|s_t))p_\theta(s_t, a_t| s_{t'}, a_{t'}) \mathrm d s_{t} \mathrm d a_{t} \Big)p_\theta(s_{t'}, a_{t'}) \mathrm d s_{t'} \mathrm d a_{t'}\\
&= \int r(s_{t'}, a_{t'}) \Big(\int \Big(\int\nabla_\theta \log(\pi_\theta(a_t|s_t)) \pi_\theta(a_t|s_t)\mathrm d a_{t}\Big)p_\theta(s_t| s_{t'}, a_{t'}) \mathrm d s_{t}  \Big)p_\theta(s_{t'}, a_{t'}) \mathrm d s_{t'} \mathrm d a_{t'}.
\end{align*}

Applying the log-derivative trick gives and using that $\pi_\theta(\cdot |s_t)$ is a probability distribution, we arrive at
$$\int\nabla_\theta \log\pi_\theta(a_t|s_t) \pi_\theta(a_t|s_t)\mathrm d a_t =\int\nabla_\theta \pi_\theta(a_t|s_t) \mathrm d a_t = \nabla_\theta\int \pi_\theta(a_t|s_t) \mathrm d a_t = \nabla_\theta1 = 0.$$

Hence, we obtain the **Reward-to-Go Policy Gradient Theorem**

\begin{align*}
\nabla_\theta J(\pi_\theta) &= \mathbb E_{\tau \sim \pi_\theta}\Big[\sum_{t < T} \nabla_\theta \log \pi_\theta(a_t|s_t)\sum_{t' \ge t} r_{t'}\Big].
\end{align*}

## Reward-to-Go for the Cartpole Example

In order to implement the reward-to-go policy gradient for the ``CartPole`` example, we only need to change the method called at the end of each episode.


```python
def finish_episode():
    """Apply optimizer to policy after each episode
    """
    
    R = 0
    policy_loss = []
    returns = []
    
    for r in policy.rewards[::-1]:
        R = r + GAMMA * R
        
        returns.insert(0, R)
    returns = torch.tensor(returns)
    
    for log_prob, R in zip(policy.saved_log_probs, 
                          returns):
        policy_loss.append(-log_prob * R)
        
    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum()
    policy_loss.backward()
    optimizer.step()
    
    del policy.rewards[:]
    del policy.saved_log_probs[:]
```

We see that this modification makes the training substantially more effective.

## Baselines

By the policy-gradient theorem, 
$$\sum_{t < T} \nabla_\theta \log \pi_\theta(a_t|s_t)\sum_{t' \ge t} r_{t'}$$
is an unbiased estimate for the gradient of the expected total rewards with respect to the policy parameters. However, as we have seen in the simulations, the variance can be substantial.

A powerful method to reduce the variance is to correct the above estimator by a **baseline** $b(s)$. That is, we now perform gradient descent with respect to expressions of the form
$$\sum_{t < T} \nabla_\theta \log \pi_\theta(a_t|s_t)\Big(\sum_{t' \ge t} r_{t'} - b(s_t)\Big).$$
As long as the baseline $b$ depends only on the state, the resulting estimator is still unbiased. Indeed, for any $t < T$,
\begin{align*}
\mathbb E_{\pi_\theta}[ \nabla_\theta \log \pi_\theta(a_t|s_t) b(s_t)] = \int \Big(\int \nabla_\theta \log \pi_\theta(a_t|s_t)   \pi_\theta(a_t|s_t)\mathrm d a_t\Big) p_\theta(s_t) b(s_t) \mathrm d s_t.
\end{align*}
and
\begin{align*}
\int \nabla_\theta \log \pi_\theta(a_t|s_t)   \pi_\theta(a_t|s_t)\mathrm d a_t=\int \nabla_\theta  \pi_\theta(a_t|s_t)  \mathrm d a_t  =   \nabla_\theta 1= 0.
\end{align*}

Since there is substantial freedom in selecting the baseline $b(s)$, this raises the question on what would be an optimal choice. In order to achieve the largest amount of variance reduction, we should define $b(s_t)$ as the conditional expectation of the reward-to-go $\sum_{t' \ge t} r_{t'}$ given $s$.  That is, we put
$$b(s_t) = V^{\pi_\theta}(s_t) = \mathbb E_{\pi_\theta}\Big[\sum_{t' \ge t} r_{t'}\,\big|\,s_t\Big]$$

## Baseline for the Cartpole Example

In practice, a crude approximation to this idea would be to subtract the empirical average instead of the conditional mean.


```python
eps = 1e-7
def finish_episode():
    """Apply optimizer to policy after each episode
    """
    
    R = 0
    policy_loss = []
    returns = []
    
    for r in policy.rewards[::-1]:
        R = r + GAMMA * R
        returns.insert(0, R)
        
    returns = torch.tensor(returns)
    returns = (returns - returns.mean()) / (returns.std() + eps)
    
    for log_prob, R in zip(policy.saved_log_probs, returns):
        policy_loss.append(-log_prob * R)
        
    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum()
    policy_loss.backward()
    optimizer.step()
    
    del policy.rewards[:]
    del policy.saved_log_probs[:]
    
```

Again, we witness a substantially more stable training in comparison to the setting without baselines.


```python
reinforce(env)
```

    10.5


## Policy Gradients with the Q Function

Finally, in some settings it also helps to apply the rewards-to-go by the $Q$-function. That is, to have 
$$\nabla_\theta J(\pi_\theta) = \mathbb E_{\tau \sim \pi_\theta}\Big[\sum_{t < T}Q^{\pi_\theta}(s_t, a_t) \nabla_\theta \log \pi_\theta(a_t|s_t)\Big]$$

In other words, we want to show that for every $t < T$,
$$\mathbb E_{\tau \sim \pi_\theta}\Big[\nabla_\theta \log \pi_\theta(a_t|s_t)\sum_{t' \ge t} r_{t'}\Big]  = \mathbb E_{\tau \sim \pi_\theta}\big[Q^{\pi_\theta}(s_t, a_t) \nabla_\theta \log \pi_\theta(a_t|s_t)\big].$$
To achieve this goal, we condition on $s_t, a_t$. Then,
\begin{align*}
\mathbb E_{\tau \sim \pi_\theta}\Big[\nabla_\theta \log \pi_\theta(a_t|s_t)\sum_{t' \ge t} r_{t'}\Big] &= \mathbb E_{\tau \sim \pi_\theta}\Big[\mathbb E_{\tau \sim \pi_\theta}\Big[\sum_{t' \ge t} r_{t'}\Big|s_t,a_t\Big]\nabla_\theta \log \pi_\theta(a_t|s_t)\Big]= \mathbb E_{\tau \sim \pi_\theta}\Big[Q^{\pi_\theta}(s_t, a_t)\nabla_\theta \log \pi_\theta(a_t|s_t)\Big], 
\end{align*}
as asserted.

Again, we can use the value function $V^{\pi_\theta}$ as a baseline and obtain the policy gradient theorem for the **advantage function $A^{\pi_\theta}(s, a) = Q^{\pi_\theta}(s, a) - V^{\pi_\theta}(s)$**. That is,
$$\nabla_\theta J(\pi_\theta) = \mathbb E_{\tau \sim \pi_\theta}\Big[\sum_{t < T}A^{\pi_\theta}(s_t, a_t) \nabla_\theta \log \pi_\theta(a_t|s_t)\Big].$$

60f2d7f1e982ecdd78ff707e90f62a0f608ee5bf
ipynb
Jupyter Notebook
notebooks/reinforce.ipynb
Christian-Hirsch/Space_Invaders
d0aedfc7d55d45c3f7cdff774facbfb0b93020f8
MIT
2019-10-09T11:37:10.000Z
2019-10-09T11:37:10.000Z
notebooks/reinforce.ipynb
Christian-Hirsch/Space_Invaders
d0aedfc7d55d45c3f7cdff774facbfb0b93020f8
MIT
notebooks/reinforce.ipynb
Christian-Hirsch/Space_Invaders
d0aedfc7d55d45c3f7cdff774facbfb0b93020f8
MIT
Qwen/Qwen-72B
1. YES
2. YES