class_docstr: ## PPO Loss
  
  Here's how the PPO update rule is derived.
  
  We want to maximize policy reward
   $$\max_    heta J(\pi_     heta) =
     \mathop{\mathbb{E}}_{    au \sim \pi_    heta}\Biggl[\sum_{t=0}^\infty \gamma^t r_t \Biggr]$$
   where $r$ is the reward, $\pi$ is the policy, $    au$ is a trajectory sampled from policy,
   and $\gamma$ is the discount factor between $[0, 1]$.
  
  egin{align}
  \mathbb{E}_{        au \sim \pi_    heta} \Biggl[
   \sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
  \Biggr] &=
  \
  \mathbb{E}_{        au \sim \pi_    heta} \Biggl[
    \sum_{t=0}^\infty \gamma^t \Bigl(
     Q^{\pi_{OLD}}(s_t, a_t) - V^{\pi_{OLD}}(s_t)
    \Bigr)
   \Biggr] &=
  \
  \mathbb{E}_{        au \sim \pi_    heta} \Biggl[
    \sum_{t=0}^\infty \gamma^t \Bigl(
     r_t + V^{\pi_{OLD}}(s_{t+1}) - V^{\pi_{OLD}}(s_t)
    \Bigr)
   \Biggr] &=
  \
  \mathbb{E}_{        au \sim \pi_    heta} \Biggl[
    \sum_{t=0}^\infty \gamma^t \Bigl(
     r_t
    \Bigr)
   \Biggr]
   - \mathbb{E}_{     au \sim \pi_    heta}
      \Biggl[V^{\pi_{OLD}}(s_0)\Biggr] &=
  J(\pi_      heta) - J(\pi_{ heta_{OLD}})
  \end{align}
  
  So,
   $$\max_    heta J(\pi_     heta) =
     \max_    heta \mathbb{E}_{       au \sim \pi_    heta} \Biggl[
        \sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
     \Biggr]$$
  
  Define discounted-future state distribution,
   $$d^\pi(s) = (1 - \gamma) \sum_{t=0}^\infty \gamma^t P(s_t = s | \pi)$$
  
  Then,
  
  egin{align}
  J(\pi_      heta) - J(\pi_{ heta_{OLD}})
  &= \mathbb{E}_{     au \sim \pi_    heta} \Biggl[
   \sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
  \Biggr]
  \
  &= rac{1}{1 - \gamma}
   \mathbb{E}_{s \sim d^{\pi_ heta}, a \sim \pi_      heta} \Bigl[
    A^{\pi_{OLD}}(s, a)
   \Bigr]
  \end{align}
  
  Importance sampling $a$ from $\pi_{ heta_{OLD}}$,
  
  egin{align}
  J(\pi_      heta) - J(\pi_{ heta_{OLD}})
  &= rac{1}{1 - \gamma}
   \mathbb{E}_{s \sim d^{\pi_ heta}, a \sim \pi_      heta} \Bigl[
    A^{\pi_{OLD}}(s, a)
   \Bigr]
  \
  &= rac{1}{1 - \gamma}
   \mathbb{E}_{s \sim d^{\pi_ heta}, a \sim \pi_{     heta_{OLD}}} \Biggl[
    rac{\pi_ heta(a|s)}{\pi_{        heta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
   \Biggr]
  \end{align}
  
  Then we assume $d^\pi_      heta(s)$ and  $d^\pi_{  heta_{OLD}}(s)$ are similar.
  The error we introduce to $J(\pi_   heta) - J(\pi_{ heta_{OLD}})$
   by this assumption is bound by the KL divergence between
   $\pi_      heta$ and $\pi_{        heta_{OLD}}$.
  [Constrained Policy Optimization](https://arxiv.org/abs/1705.10528)
   shows the proof of this. I haven't read it.
  
  
  egin{align}
  J(\pi_      heta) - J(\pi_{ heta_{OLD}})
  &= rac{1}{1 - \gamma}
   \mathop{\mathbb{E}}_{s \sim d^{\pi_        heta} top a \sim \pi_{ heta_{OLD}}} \Biggl[
    rac{\pi_ heta(a|s)}{\pi_{        heta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
   \Biggr]
  \
  &pprox rac{1}{1 - \gamma}
   \mathop{\mathbb{E}}_{      extcolor{orange}{s \sim d^{\pi_{        heta_{OLD}}}}
   top a \sim \pi_{  heta_{OLD}}} \Biggl[
    rac{\pi_ heta(a|s)}{\pi_{        heta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
   \Biggr]
  \
  &= rac{1}{1 - \gamma} \mathcal{L}^{CPI}
  \end{align}
code: def __init__(self):
      super().__init__()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/rl/ppo/__init__.py
imports: import torch
  from labml_helpers.module import Module
  from labml_nn.rl.ppo.gae import GAE
label_desc: Use image node_id 1 to create a new ClippedPPOLoss object from inherited base classes: Module with example: obj = ClippedPPOLoss()