From af1203b9df727800332d6b00f093b615357d7810 Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Fri, 26 Jun 2020 09:52:22 +0200
Subject: [PATCH] [RLlib] Issue 8507 (PyTorch does not support custom loss).
 (#9142)

---
 rllib/examples/models/custom_loss_model.py | 37 ++++++++++++++++++----
 rllib/offline/json_reader.py               |  2 +-
 rllib/policy/torch_policy.py               |  3 ++
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/rllib/examples/models/custom_loss_model.py b/rllib/examples/models/custom_loss_model.py
index 2d3acd1de..0a536c7d2 100644
--- a/rllib/examples/models/custom_loss_model.py
+++ b/rllib/examples/models/custom_loss_model.py
@@ -121,6 +121,8 @@ class TorchCustomLossModel(TorchModelV2, nn.Module):
         nn.Module.__init__(self)
 
         self.input_files = input_files
+        # Create a new input reader per worker.
+        self.reader = JsonReader(self.input_files)
         self.fcnet = TorchFC(
             self.obs_space,
             self.action_space,
@@ -135,13 +137,30 @@ class TorchCustomLossModel(TorchModelV2, nn.Module):
 
     @override(ModelV2)
     def custom_loss(self, policy_loss, loss_inputs):
-        # Create a new input reader per worker.
-        reader = JsonReader(self.input_files)
-        input_ops = reader.tf_input_ops()
+        """Calculates a custom loss on top of the given policy_loss(es).
+
+        Args:
+            policy_loss (List[TensorType]): The list of already calculated
+                policy losses (as many as there are optimizers).
+            loss_inputs (TensorStruct): Struct of np.ndarrays holding the
+                entire train batch.
+
+        Returns:
+            List[TensorType]: The altered list of policy losses. In case the
+                custom loss should have its own optimizer, make sure the
+                returned list is one larger than the incoming policy_loss list.
+                In case you simply want to mix in the custom loss into the
+                already calculated policy losses, return a list of altered
+                policy losses (as done in this example below).
+        """
+        # Get the next batch from our input files.
+        batch = self.reader.next()
 
         # Define a secondary loss by building a graph copy with weight sharing.
         obs = restore_original_dimensions(
-            tf.cast(input_ops["obs"], tf.float32), self.obs_space)
+            torch.from_numpy(batch["obs"]).float(),
+            self.obs_space,
+            tensorlib="torch")
         logits, _ = self.forward({"obs": obs}, [], None)
 
         # You can also add self-supervised losses easily by referencing tensors
@@ -155,11 +174,15 @@ class TorchCustomLossModel(TorchModelV2, nn.Module):
         action_dist = TorchCategorical(logits, self.model_config)
         self.policy_loss = policy_loss
         self.imitation_loss = torch.mean(
-            -action_dist.logp(input_ops["actions"]))
-        return policy_loss + 10 * self.imitation_loss
+            -action_dist.logp(torch.from_numpy(batch["actions"])))
+
+        # Add the imitation loss to each already calculated policy loss term.
+        # Alternatively (if custom loss has its own optimizer):
+        # return policy_loss + [10 * self.imitation_loss]
+        return [l + 10 * self.imitation_loss for l in policy_loss]
 
     def custom_stats(self):
         return {
-            "policy_loss": self.policy_loss,
+            "policy_loss": torch.mean(self.policy_loss),
             "imitation_loss": self.imitation_loss,
         }
diff --git a/rllib/offline/json_reader.py b/rllib/offline/json_reader.py
index c2f0ac6ba..b72d139cc 100644
--- a/rllib/offline/json_reader.py
+++ b/rllib/offline/json_reader.py
@@ -45,7 +45,7 @@ class JsonReader(InputReader):
                 logger.warning(
                     "Treating input directory as glob pattern: {}".format(
                         inputs))
-            if urlparse(inputs).scheme:
+            if urlparse(inputs).scheme not in ["d", ""]:
                 raise ValueError(
                     "Don't know how to glob over `{}`, ".format(inputs) +
                     "please specify a list of files to read instead.")
diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py
index eac18200d..bd701ca0b 100644
--- a/rllib/policy/torch_policy.py
+++ b/rllib/policy/torch_policy.py
@@ -240,6 +240,9 @@ class TorchPolicy(Policy):
         train_batch = self._lazy_tensor_dict(postprocessed_batch)
         loss_out = force_list(
             self._loss(self, self.model, self.dist_class, train_batch))
+        # Call Model's custom-loss with Policy loss outputs and train_batch.
+        if self.model:
+            loss_out = self.model.custom_loss(loss_out, train_batch)
         assert len(loss_out) == len(self._optimizers)
         # assert not any(torch.isnan(l) for l in loss_out)