Reinforcement Learning

The Networks

There are two networks required for PPO training.

The “Policy Network”, which we’ll denote with the function $\pi(\cdot)$, is a mapping from the state space to a distribution over actions.

Mathematically, $\pi: \pmb s \to \pmb a$ where $\pmb s \in \mathbb R^s$ and $\pmb a \in \mathbb R^a$ where $s$ is the dimension of the state space and $a$ is the dimension of the action space. Note that since $\pi(\cdot)$ maps to the state space, $a_i \geq 0$ for all $i$ and $\sum_{i=1}^p a_i = 1$.

The “Value Network”, which we’ll denote with the function $V(\cdot)$, is a mapping from the state space to “rewards” over the action space.

Concisely, $V: \pmb s \to r$ where $\pmb s$ is the same as defined before and $r \in R$. Notice that unlike $\pi$, $V$ is scalar valued. It’s a measure of the desirability of the given state.

Simulated Environment

Information about $\pi(\cdot)$ and $V(\cdot)$ needs to be revealed somehow. In reinforcement learning, this information comes from the “environment” which we’ll represent with the function $E(\cdot)$. The environment is responsible for knowing how to 1) generate an initial state 2) update the state based on an action 3) provide an award associated with an action 4) know when the game is over.

A classic environment where reinforcement learning can be used is the “cart pole” simulation (also known as an inverted pendulum to physicists). The game is simple; there is a pole on a cart and the objective is the keep the pole in the air. We’ll do this experiment in a 2-D world where the cart can either be moved left or right ($p=2$). Although there are only 2 directions our cart can move, we’ll need to choose a single direction to move the cart at each point in time. We’ll call our choice $c \in \{0,1\}$ where $0$ represents moving left and $1$ represent moving right.

The state of the cart pole system can be summarized by a vector of 4 values: position of the cart, velocity of the cart, angle of the pole, and rotation of the pole ($s=4$).

Now that we know the parts of our system we can describe $E(\cdot)$ as being a list-valued function $E(\pmb s_t, c_t) = [\pmb s_{t+1}, r_t, d_{t+1}]$. Here, $r_t$ is known as the “reward” which is the reward associated with the action (choice) $c_t$. The element, $d_{t+1}$, is an indicator variable for if the game is over.

In the cart pole scenario, once the pole has fallen over the game is over. When training a netork, nothing after this point should be learned because no action has an effect on the end result from this point onward.

The Goal

We’ve been introduced to our three main components of reinforcement learning: the policy $\pi(\cdot)$, The value $V(\cdot)$, and the environment $E(\cdot)$. Though $E(\cdot)$ must be known before training time, nothing is known about $\pi(\cdot)$ and $V(\cdot)$ so we initialize networks $\hat \pi(\cdot)$ and $\hat V(\cdot)$. The initial versions of this network we’ll denote $\hat \pi^0(\cdot)$ and $\hat V^0(\cdot)$ where the superscript $0$ indicates that no training has taken place. If all goes well, as more training takes place $\hat \pi(\cdot)$ and $\hat V(\cdot)$ will converge to $\pi(\cdot)$ and $V(\cdot)$ respectively. That is to say, as long as our networks have a good architecture then $\lim\limits_{n \to \infty} \hat \pi^n = \pi$ and $\lim\limits_{n \to \infty} \hat V^n = V$.

Learing Via Simulation

Though $\hat \pi$ and $\hat V$ start off no better than random number generators, through interacting with $E(\cdot)$ they’ll become more like $\pi$ and $V$. First, we run a full simulation to see how we’re doing.

Generating an experience

We start off with some initial state $\pmb s_0$ (from our environment) and we want to let our network choose what to do given this state. Therefore we compute $\pmb a_0 = \hat \pi^0(\pmb s_0)$ where $\pmb a_0$ is a distribution of choices over our discrete action space. We sample from $\pmb a_0$ by treating it as a discrete distribution where the probability of selecting each action is given by $a_i$.

After sampling the action, we’ll have scalar $c_0 \in \{1,2,\cdots, a\}$ which is our network’s response to $\pmb s_0$.

It is key to note that since we are sampling from our distribution, there is an element of randomness in our model’s deciding which path to take.

We now compute $E(\pmb s_t, c_t) = [\pmb s_{t+1}, r_t, d_{t+1}]$. We now have reward $r_t$ associated with action $c_t$ given under $\pmb a_t$ which was generated by $\hat \pi^0(\cdot)$.

Given the new $s_{t+1}$ and assuming $d_{t+1} \neq 1$ we repeat the process some $T$ times.

In pseudocode we have:

s = E.start();

game_over = False;
steps = 0;
experience = set();

while (not game_over and steps < T) {
     c, pmf = policy(s);

     next_s, r, game_over = E(s, c);

     steps += 1;
     experience.add(s, c, r, pmf);

     s = next_s;
};

Correcting rewards

Though each action has a reward associated with it, some actions with lower initial rewards by lead to higher rewards in the future.

Therefore the rewards with each action aren’t actually the raw rewards that we’ve collected. We update these rewards with information about future rewards via:

$$\begin{align*} \tilde r_{t} = \sum_{i=t}^T \gamma^{i-t} r_{i} \tag{adjusted reward} \end{align*}$$

Where $0 < \gamma < 1$ and called the “discount rate” or “patience factor”. A reward today is worth more than a reward tomorrow. If $\gamma = 1$ then a reward tomorrow is just as good as reward today and if $\gamma = 0$ then a reward tomorrow is worth nothing today.

Now that the rewards have been updated to more reasonable measurements, our network is ready to learn from its experience.

Learning

Now we’ll use our set of experiences, call it $e$, to learn. This set looks like this $e^0 = \{(\pmb s_0, c_0, r_0, \pmb a_0), (\pmb s_1, c_1, r_1, \pmb a_1), \dots, (\pmb s_T, c_T, r_T, \pmb a_T)\}$. We’ll sample from these experiences to learn.

There’s two losses that will make up the total loss: value loss and policy loss

Value Loss

Value loss is a straightforward intuitive calculation. It’s the mean squared error (MSE) between the predicted rewards by the value network and the returns obtained from experience.

$$\text{ValueLoss} = MSE(\hat {\pmb r}, \tilde {\pmb r})$$

Where $\hat {\pmb r} = V(\pmb s)$ and $\tilde {\pmb r}$ is the adjusted reward.

Policy Loss

Policy loss is bit more involved. The equation is

$$\begin{align*} \text{PolicyLoss} = -&\text{Mean}\left(\min\left[\frac{\pi'(a|\pmb s)}{\pi(a|\pmb s)}\hat A, \text{clip}\left(\frac{\pi'(a|\pmb s)}{\pi(a|\pmb s)}, 1- \epsilon, 1 + \epsilon \right)\hat A \right]\ \right) \end{align*}$$

Where $\pi'$ is what would happen if we accepted the random walk as the new policy and $\pi$ is the current (old) policy. $\text{clip}(\cdot)$ is a function that uses $\epsilon$ to determine how large of a step we’re comfortable iwth making in the optimization space. Finally $\hat A$ is defined as

$$\begin{align*} \hat A &= \tilde{r} - \hat{r} \\ &= V'(a |\pmb s) - V(a | \pmb s) \\ &= V_{\text{new}}(a | \pmb s) - V_{\text{current}}(a | \pmb s) \\ &= r_{\text{experience}} - r_{\text{predicted}} \end{align*}$$

And is referred to as “advantage”. Advantage is positive when tilting the distribution toward the explored route would increase rewards compared to the decision under the current policy. Since capitalizing on advantage is good, that’s why PolicyLoss has a negative sign on the outside of mean.

The term $\pi'/\pi$ is called the “policy ratio” which determines how much the network follows the advantage signal. When $\pi'/\pi > 1$ then the advatange signal is amplified. This is desirable because as an action’s relative weighting increases, the importance of accurately gauging the reward associated with that action also increases which is precisely what this policy ratio reflects.

The policy ratio will be $1$ before updating the function $\pi$ because initially we have $\pi^{\prime\ (0)}/\pi^{\prime\ (0)}$ however after backpropping on the the total loss (which is equal to $\text{ValueLoss} + \text{PolicyLoss}$) we’ll be computing the $\pi^{\prime\ (1)}/\pi^{\prime\ (0)}$ after sampling the next batch from memories. We continually update $\pi^{\prime\ (i)} \to \pi^{\prime\ (i+1)}$ by sampling new batches from the given set of memories until we decide to get a obtain a new set of memories.

The minimum loss term is chosen from the two policy terms for optimization stability.

In pseudo code this training process looks like

function learn_from_memory(policy, value, memories) {
    for (number_of_epochs) {
        for mem in memories {
            s, c, r, pmf = mem;

            hat_r = value(s);
            value_loss = MSE(hat_r, r);


            proba_c = policy(s).at(c);
            old_proba_c = pmf.at(c);
            change = proba_c / old_proba_c;

            adv = r - hat_r;

            policy_loss = min([
                change * adv,
                clip(change, 1-eps, 1+eps) * adv
            ]);

            tot_loss = value_loss + policy_loss;

            back_prop_step(tot_loss);
        };
    };
};

When this point is reached, we repeat the process outlined except now we have $\hat \pi^1$ where $1$ denotes that the model has learned from one set of memories. Each successive set of memories should be more advanced than the previous. After enough training, if the architecture is good $\hat \pi \to \pi$ and $\hat V \to V$ and our model now understands the game it’s playing.