Add KL divergence cap

Author: Omegastick

include/cpprl/algorithms/ppo.h

@@ -16,7 +16,7 @@ class PPO : public Algorithm
 {
   private:
     Policy &policy;
-    float actor_loss_coef, value_loss_coef, entropy_coef, max_grad_norm, original_learning_rate, original_clip_param;
+    float actor_loss_coef, value_loss_coef, entropy_coef, max_grad_norm, original_learning_rate, original_clip_param, kl_target;
     int num_epoch, num_mini_batch;
     std::unique_ptr<torch::optim::Adam> optimizer;
 
@@ -30,7 +30,8 @@ class PPO : public Algorithm
         float entropy_coef,
         float learning_rate,
         float epsilon = 1e-8,
-        float max_grad_norm = 0.5);
+        float max_grad_norm = 0.5,
+        float kl_target = 0.01);
 
     std::vector<UpdateDatum> update(RolloutStorage &rollouts, float decay_level = 1);
 };

src/algorithms/ppo.cpp

@@ -24,14 +24,16 @@ PPO::PPO(Policy &policy,
          float entropy_coef,
          float learning_rate,
          float epsilon,
-         float max_grad_norm)
+         float max_grad_norm,
+         float kl_target)
     : policy(policy),
       actor_loss_coef(actor_loss_coef),
       value_loss_coef(value_loss_coef),
       entropy_coef(entropy_coef),
       max_grad_norm(max_grad_norm),
       original_learning_rate(learning_rate),
       original_clip_param(clip_param),
+      kl_target(kl_target),
       num_epoch(num_epoch),
       num_mini_batch(num_mini_batch),
       optimizer(std::make_unique<torch::optim::Adam>(
@@ -57,6 +59,9 @@ std::vector<UpdateDatum> PPO::update(RolloutStorage &rollouts, float decay_level
     float total_value_loss = 0;
     float total_action_loss = 0;
     float total_entropy = 0;
+    float kl_divergence = 0;
+    float kl_early_stopped = -1;
+    int num_updates = 0;
 
     // Epoch loop
     for (int epoch = 0; epoch < num_epoch; ++epoch)
@@ -86,6 +91,17 @@ std::vector<UpdateDatum> PPO::update(RolloutStorage &rollouts, float decay_level
                 mini_batch.masks,
                 mini_batch.actions);
 
+            // Calculate approximate KL divergence for info and early stopping
+            kl_divergence = (mini_batch.action_log_probs - evaluate_result[1])
+                                .mean()
+                                .item()
+                                .toFloat();
+            if (kl_divergence > kl_target * 1.5)
+            {
+                kl_early_stopped = num_updates;
+                goto finish_update;
+            }
+
             // Calculate difference ratio between old and new action probabilites
             auto ratio = torch::exp(evaluate_result[1] -
                                     mini_batch.action_log_probs);
@@ -114,22 +130,34 @@ std::vector<UpdateDatum> PPO::update(RolloutStorage &rollouts, float decay_level
             loss.backward();
             // TODO: Implement gradient norm clipping
             optimizer->step();
+            num_updates++;
 
             total_value_loss += value_loss.item().toFloat();
             total_action_loss += action_loss.item().toFloat();
             total_entropy += evaluate_result[2].item().toFloat();
         }
     }
 
-    auto num_updates = num_epoch * num_mini_batch;
-
+finish_update:
     total_value_loss /= num_updates;
     total_action_loss /= num_updates;
     total_entropy /= num_updates;
 
-    return {{"Value loss", total_value_loss},
-            {"Action loss", total_action_loss},
-            {"Entropy", total_entropy}};
+    if (kl_early_stopped > -1)
+    {
+        return {{"Value loss", total_value_loss},
+                {"Action loss", total_action_loss},
+                {"Entropy", total_entropy},
+                {"KL divergence", kl_divergence},
+                {"KL divergence early stop update", kl_early_stopped}};
+    }
+    else
+    {
+        return {{"Value loss", total_value_loss},
+                {"Action loss", total_action_loss},
+                {"Entropy", total_entropy},
+                {"KL divergence", kl_divergence}};
+    }
 }
 
 TEST_CASE("PPO")