fix: KL runs by calculating KL (BC) loss only when above skipped level

Author: MichalBortkiewicz

experiment_code/hackrl/experiment.py

@@ -611,14 +611,23 @@ def compute_entropy_loss(logits, stats=None):
     return -torch.mean(entropy_per_timestep)
 
 
-def compute_kickstarting_loss(student_logits, expert_logits):
+def compute_kickstarting_loss(student_logits, expert_logits, mask: torch.Tensor):
     T, B, *_ = student_logits.shape
-    return torch.nn.functional.kl_div(
+    if not mask:
+        return torch.nn.functional.kl_div(
+            F.log_softmax(student_logits.reshape(T * B, -1), dim=-1),
+            F.log_softmax(expert_logits.reshape(T * B, -1), dim=-1),
+            log_target=True,
+            reduction="batchmean",
+        )
+    loss = torch.nn.functional.kl_div(
         F.log_softmax(student_logits.reshape(T * B, -1), dim=-1),
         F.log_softmax(expert_logits.reshape(T * B, -1), dim=-1),
         log_target=True,
-        reduction="batchmean",
+        reduction="none",
     )
+    loss = loss.T * mask
+    return loss.sum() / B / T
 
 
 def compute_policy_gradient_loss(
@@ -888,6 +897,8 @@ def compute_gradients(data, sleep_data, learner_state, stats):
         stats["inverse_loss"] += inverse_loss.item()
 
     if FLAGS.use_kickstarting:
+        # TODO phase 2: add regularization only mask, when we reach a particular lvl
+
         kickstarting_loss = FLAGS.kickstarting_loss * compute_kickstarting_loss(
             learner_outputs["policy_logits"],
             actor_outputs["kick_policy_logits"],
@@ -912,6 +923,7 @@ def compute_gradients(data, sleep_data, learner_state, stats):
         kickstarting_loss_bc = FLAGS.kickstarting_loss_bc * compute_kickstarting_loss(
             ttyrec_predictions["policy_logits"],
             ttyrec_predictions["kick_policy_logits"],
+            torch.flatten(ttyrec_data["mask"], 0, 1).int()
         )
         FLAGS.kickstarting_loss_bc *= FLAGS.kickstarting_decay_bc
         total_loss += kickstarting_loss_bc

experiment_code/mrunner_exps/skipping_levels/monk-APPO-AA-KL-T-skip.py renamed to experiment_code/mrunner_exps/skipping_levels/monk-APPO-AA-KL-T-skip-proper.py

@@ -28,7 +28,7 @@
 params_grid = [
     {
         "exp_tags": [f"{name}-4"],
-        "seed": list(range(2, 3)),
+        "seed": list(range(5)),
         # load from checkpoint
         "unfreeze_actor_steps": [0],
         "use_checkpoint_actor": [True],
@@ -42,7 +42,7 @@
     },
     {
         "exp_tags": [f"{name}-3"],
-        "seed": list(range(3)),
+        "seed": list(range(5)),
         # load from checkpoint
         "unfreeze_actor_steps": [0],
         "use_checkpoint_actor": [True],
@@ -56,7 +56,7 @@
     },
     {
         "exp_tags": [f"{name}-2"],
-        "seed": list(range(3)),
+        "seed": list(range(5)),
         # load from checkpoint
         "unfreeze_actor_steps": [0],
         "use_checkpoint_actor": [True],
@@ -70,7 +70,7 @@
     },
     {
         "exp_tags": [f"{name}-1"],
-        "seed": list(range(3)),
+        "seed": list(range(5)),
         # load from checkpoint
         "unfreeze_actor_steps": [0],
         "use_checkpoint_actor": [True],
@@ -84,7 +84,7 @@
     },
     {
         "exp_tags": [f"{name}-0"],
-        "seed": list(range(3)),
+        "seed": list(range(5)),
         # load from checkpoint
         "unfreeze_actor_steps": [0],
         "use_checkpoint_actor": [True],