Reinforcement Learning (RL) is a framework to train an intelligent agent to take a sequence of actions in order to maximize a reward, given the feedback from a dynamic environment.

flowchart LR
%% Define the nodes
Agent[Agent]
Environment[Environment]

%% Define the interactions (edges)
Agent -- Action (A1, A2, ..., At) --> Environment
Environment -- State (S1, S2, ..., St) --> Agent
Environment -- Reward (R1, R2, ..., Rt) --> Agent

Gaming Agent

One of the most common use cases is gaming, where the agent is trained to win the game by controlling te keypads based on the gaming environment. For example, in the game of Pong, an agent can perform two actions $A_t$ at every timestamp: UP and DOWN. The environment will present the visual outcome as state $S_t$, which is a video frame of two paddles and a ball. In addition, the environment will reward the model (e.g. $R_t=1$) when the agent wins the game, and otherwise penalize it (e.g. $R_t=-1$).

Discounted Reward

Practically, the reward (or penalty) will be carried over to the previous timestamps with a discounting factor $\gamma$:

def discount_rewards(rewards, gamma=0.99):
  discounted = np.zeros_like(rewards, dtype=np.float32)
  running_add = 0
  for t in reversed(range(len(rewards))):
    if rewards[t] != 0:
      # BOUNDARY: Reset the chain when a point is scored
      running_add = 0
    running_add = running_add * gamma + rewards[t]
    discounted[t] = running_add
  return discounted

Policy Network

The agent learns a policy network $\pi_\theta(a|s)$ to decide what actions ($a$) to take given the current state ($s$). It produces a probablistic distribution of actions to sample from at each timestamp, parameterized by $\theta$.

Reward Optimization

The objective is to maximize the expected reward:

To optimize the reward, we compute its gradient:

Loss Function

To make the policy network trainable, we define a loss function that produces the same gradient (but negative) as above:

As we sample from the dataset, we are essentially trending towards the expected value.

Advantage vs Reward

However, this loss function will produce very high variance, mostly due to the noisy reward signal: some intermediate actions may be suboptimal, but the same reward will be applied to all past actions.

To mitigate the issue, people often replace reward $R$ with advantage $A$ that has more stable mean and variance. There are many design choices for this advantage. One simple option is normalization: we can define advantage as the normalized reward within the batch.

Another commonly used option is to use a critic network $Q_\phi(s, a)$ and $V_\phi(s)$, where $Q_\phi(s, a)$ predicts the reward given the current $s$ and $a$, and $V_\phi(s)$ predicts the expected reward given the current $s$ only.

As the advantage $A$ is not a function of $\theta$, the loss $L$ and its gradient remain the same.

Training Loop

The training loop consists of 4 steps:

1. Data Collection: The agent plays the game with the frozen policy $\pi_\theta(a|s)$ until certain criteria are met, e.g. M frames are seen or N games are finished.
2. Optimization: We calculate the advantage $A$ and the loss $L$, and propagate the gradient to the policy network $\pi_\theta(a|s)$. We run this optimization for a few steps.
3. Policy Update: We update the weight $\theta$ of the policy network $\pi_\theta(a|s)$.
4. Repeat: We repeat the process.

for _ in range(num_iterations):
  # 1. Data collection
  states, actions, rewards, old_probs = collect_rollout(
      env, policy_net
  )
  # Calculate advantage using the critic
  advantages = compute_advantages(rewards, critic_net, states)
  
  # 2. Optimization
  optimizer.zero_grad()
  for _ in range(num_steps):
    # 3. Policy update
    new_probs = policy_net(states)
    loss = compute_loss(new_probs, old_probs, advantages)
    loss.backward()
    optimizer.step()

Proximal Policy Optimization (PPO)

To further reduce the variance, Proximal Policy Optimization (PPO) defines a bounded loss function:

When $A>0$, it is a good action, and otherwise a bad action. When $A$ and $r-1$ has the same sign, the policy moves toward the same direction, and otherwise a different direction. Here is a case-by-case study of how this bounded loss function plays its role:

As is shown from the above table, there is no gradient to $\theta$ (∂L/∂θ = 0) when the new policy moves towards the same direction as the reward, and the change is too much:

• $A\ge0$, $r>1+\epsilon$: Action is good, but the new policy increases the probability too much.
• $A<0$, $r<1-\epsilon$: Action is bad, but the new policy decreases the probability too much.

Reference

• https://karpathy.github.io/2016/05/31/rl/