Repository: ray-project/ray

Issue:
[rllib] Pytorch missing custom loss
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### What is the problem?
Running `rllib/examples/custom_loss.py` example, the tensorflow implementation displays the custom loss, but Pytorch does not.  During the process, I had to update the `rllib/examples/models/custom_loss_model.py` models method name from `custom_stats `to `metrics`.

*Ray version and other system information (Python version, TensorFlow version, OS):*
Ray version: latest wheel
Python: 3.6.8
TF: 2.1
OS: RHEL 7.7

### Reproduction (REQUIRED)
Please provide a script that can be run to reproduce the issue. The script should have **no external library dependencies** (i.e., use fake or mock data / environments):

Use existing example script in Rllib

If we cannot run your script, we cannot fix your issue.

- [x ] I have verified my script runs in a clean environment and reproduces the issue.
- [x ] I have verified the issue also occurs with the [latest wheels](https://docs.ray.io/en/latest/installation.html).



Discussion:
#8193 
Sorry, this may have slipped through the cracks. I'll take a look and provide the missing example.
Got it, `ModelV2.custom_loss()` is not supported for PyTorch yet at all (ignored). Will fix this.

Yes, `metrics` is the new name (ModelV2) for the old `custom_stats` (Model(V1)).
I'm looking closer at the current tf policies, and I'm actually struggling to see where `model.custom_loss` has been implemented in tf IMPALA either... Can you highlight the part of the code where this is done for my benefit? In the case of IMPALA, I know it calls the `model.custom_loss` method, as I can see the custom loss metrics I've produced, but what is not clear is that it is actually making it into the gradient calculations... `build_vtrace_loss` seems to just call  the class `VTraceLoss`, where I don't see any references to the output of` model.custom_loss`. and dynamic_tf_policy seems to just return the output of this loss function... so I'm struggling to see where it comes into play.
Looks like it is only used once [here](https://github.com/ray-project/ray/blob/master/rllib/policy/tf_policy.py#L215) and from there on it's accessed through `self._loss`.
Thanks for that, not sure how that eluded me… it looks like on further review my issue on the TF side was using 1.x style tf.losses.get_regularization_loss rather than model.losses … once I switched everything is working for TF models.


On May 28, 2020, at 4:50 PM, Jan Blumenkamp <notifications@github.com> wrote:

> 
> Looks like it is only used once here and from there on it's accessed through self._loss.
> 
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub, or unsubscribe.
> 


> Got it, `ModelV2.custom_loss()` is not supported for PyTorch yet at all (ignored). Will fix this.
> 
> Yes, `metrics` is the new name (ModelV2) for the old `custom_stats` (Model(V1)).

Hi! Is there any update on this? If update is not expected soon, I might give it a shot with tensorflow side. Many thanks.
Hey, @cedros23 . Taking another look right now.

Resolving patch:
diff --git a/rllib/examples/models/custom_loss_model.py b/rllib/examples/models/custom_loss_model.py
--- a/rllib/examples/models/custom_loss_model.py
+++ b/rllib/examples/models/custom_loss_model.py
@@ -121,6 +121,8 @@ def __init__(self, obs_space, action_space, num_outputs, model_config,
         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 @@ def forward(self, input_dict, state, seq_lens):
 
     @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 @@ def custom_loss(self, policy_loss, loss_inputs):
         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
--- a/rllib/offline/json_reader.py
+++ b/rllib/offline/json_reader.py
@@ -45,7 +45,7 @@ def __init__(self, inputs, ioctx=None):
                 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
--- a/rllib/policy/torch_policy.py
+++ b/rllib/policy/torch_policy.py
@@ -240,6 +240,9 @@ def learn_on_batch(self, postprocessed_batch):
         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)
 


Test patch: