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 forward(
      self,
      log_pi: torch.Tensor,
      sampled_log_pi: torch.Tensor,
      advantage: torch.Tensor,
      clip: float,
  ) -> torch.Tensor:
      # ratio $r_t(\theta) = \frac{\pi_\theta (a_t|s_t)}{\pi_{\theta_{OLD}} (a_t|s_t)}$;
      # *this is different from rewards* $r_t$.
      ratio = torch.exp(log_pi - sampled_log_pi)
  
      # ### Cliping the policy ratio
      #
      # \begin{align}
      # \mathcal{L}^{CLIP}(\theta) =
      #  \mathbb{E}_{a_t, s_t \sim \pi_{\theta{OLD}}} \biggl[
      #    min \Bigl(r_t(\theta) \bar{A_t},
      #              clip \bigl(
      #               r_t(\theta), 1 - \epsilon, 1 + \epsilon
      #              \bigr) \bar{A_t}
      #    \Bigr)
      #  \biggr]
      # \end{align}
      #
      # The ratio is clipped to be close to 1.
      # We take the minimum so that the gradient will only pull
      # $\pi_\theta$ towards $\pi_{\theta_{OLD}}$ if the ratio is
      # not between $1 - \epsilon$ and $1 + \epsilon$.
      # This keeps the KL divergence between $\pi_\theta$
      #  and $\pi_{\theta_{OLD}}$ constrained.
      # Large deviation can cause performance collapse;
      #  where the policy performance drops and doesn't recover because
      #  we are sampling from a bad policy.
      #
      # Using the normalized advantage
      #  $\bar{A_t} = \frac{\hat{A_t} - \mu(\hat{A_t})}{\sigma(\hat{A_t})}$
      #  introduces a bias to the policy gradient estimator,
      #  but it reduces variance a lot.
      clipped_ratio = ratio.clamp(min=1.0 - clip, max=1.0 + clip)
      policy_reward = torch.min(
          ratio * advantage, clipped_ratio * advantage
      )
  
      self.clip_fraction = (
          (abs((ratio - 1.0)) > clip).to(torch.float).mean()
      )
  
      return -policy_reward.mean()
  
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 2 for calling the ClippedPPOLoss obj's underlying member method code with example usage: obj.forward(log_pi, sampled_log_pi, advantage, clip) and returns: unknown