# KL divergence coefficient update doesn't make sense in RLlib's PPO implementation

I am using RLlib (Ray 1.4.0)'s implementation of PPO for a multi-agent scenario with continuous actions, and I find that the loss includes the KL divergence penalty term, apart from the surrogate loss, value function loss, and entropy.

The KL coefficient is updated in the update_kl() function as follows:

    if sampled_kl > 2.0 * self.kl_target:
self.kl_coeff_val *= 1.5
# Decrease.
elif sampled_kl < 0.5 * self.kl_target:
self.kl_coeff_val *= 0.5
# No change.
else:
return self.kl_coeff_val


I don't understand the reasoning behind this. If the point of the KL "target" is to reach the target, then why do the conditions above imply that the KL coefficient is getting larger (multiplied by 1.5 when the sampled KL is already found to be larger than the target?) when it is supposed to be made smaller instead? I feel like I am missing something here, but I am not able to get my head around it.

I would appreciate any insights on this. Thank you.

The loss in PPO has a penalty

$$-\text{kl_coeff_val}\cdot\mathrm{KL}(\pi^{\text{NN}}||\pi^{\text{Target}}).$$

Increasing this coefficient means a lower reward, so it incentivizes a smaller KL divergence. Likewise, decreasing it allows a higher KL divergence.

Please note that the official paper introduces two different algorithms for proximal policy optimization: PPO-CLIP (Sec. 3) and PPO-Penalty (Sec. 4)

The algorithm that you are describing is PPO-Penalty and the objective function is given by the following equation:

$$L^{KL} = \mathbb{E} \bigg[ \frac{\pi_\theta (a|s)}{\pi_{\theta_{old}} (a|s)} A(a,s) - \beta \text{KL}(\pi_\theta, \pi_{\theta_{old}}) \bigg]$$

Here $$\beta$$ is not a fixed hyperparameter, but is a Lagrangian multiplier for the constraint. In this case the constraint is that $$\text{KL}(\pi_\theta, \pi_{\theta_{old}}) \leq KL_{tgt}$$.

During optimization you want to increase the multiplier if the constraint is violated, and lower it otherwise.

Why do we want to do that?

If your constraint is violated too much, you want to increase the multiplier $$\beta$$ so that the next update step will pay more attention to this term and will try to take smaller update steps on the policy network. If the constraint is not violated, then you decrease the multiplier $$\beta$$ so that we can take larger update steps on the policy network.

In theory you should converge to an optimal coefficient $$\beta$$ such that you take the largest update step on the policy network, such that the constraint is not violated. In this case the constraint will be fulfilled exactly.