From ba030482542cb670f5169b0d13254ec68c927deb Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Sat, 2 Mar 2019 22:57:51 -0800
Subject: [PATCH] [rllib] TF model custom_loss() should actually allow access
 to full rollout data (#4220)

---
 doc/source/rllib-models.rst                    |  3 ++-
 .../ray/rllib/agents/ddpg/ddpg_policy_graph.py | 18 +++++++++++-------
 python/ray/rllib/evaluation/tf_policy_graph.py | 13 +++++++------
 python/ray/rllib/examples/custom_loss.py       |  6 ++++--
 python/ray/rllib/models/model.py               |  3 ++-
 python/ray/rllib/offline/input_reader.py       |  2 +-
 6 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/doc/source/rllib-models.rst b/doc/source/rllib-models.rst
index 54b4613b1..d0c1a1082 100644
--- a/doc/source/rllib-models.rst
+++ b/doc/source/rllib-models.rst
@@ -90,7 +90,7 @@ Custom TF models should subclass the common RLlib `model class <https://github.c
             return tf.reshape(
                 linear(self.last_layer, 1, "value", normc_initializer(1.0)), [-1])
 
-        def custom_loss(self, policy_loss):
+        def custom_loss(self, policy_loss, loss_inputs):
             """Override to customize the loss function used to optimize this model.
 
             This can be used to incorporate self-supervised losses (by defining
@@ -102,6 +102,7 @@ Custom TF models should subclass the common RLlib `model class <https://github.c
 
             Arguments:
                 policy_loss (Tensor): scalar policy loss from the policy graph.
+                loss_inputs (dict): map of input placeholders for rollout data.
 
             Returns:
                 Scalar tensor for the customized loss for this model.
diff --git a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
index 0089ed6bf..439671e93 100644
--- a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
+++ b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
@@ -333,13 +333,6 @@ class DDPGPolicyGraph(TFPolicyGraph):
                         self.loss.critic_loss += (
                             config["l2_reg"] * 0.5 * tf.nn.l2_loss(var))
 
-        # Model self-supervised losses
-        self.loss.actor_loss = self.p_model.custom_loss(self.loss.actor_loss)
-        self.loss.critic_loss = self.q_model.custom_loss(self.loss.critic_loss)
-        if self.config["twin_q"]:
-            self.loss.critic_loss = self.twin_q_model.custom_loss(
-                self.loss.critic_loss)
-
         # update_target_fn will be called periodically to copy Q network to
         # target Q network
         self.tau_value = config.get("tau")
@@ -375,6 +368,17 @@ class DDPGPolicyGraph(TFPolicyGraph):
             ("dones", self.done_mask),
             ("weights", self.importance_weights),
         ]
+        input_dict = dict(self.loss_inputs)
+
+        # Model self-supervised losses
+        self.loss.actor_loss = self.p_model.custom_loss(
+            self.loss.actor_loss, input_dict)
+        self.loss.critic_loss = self.q_model.custom_loss(
+            self.loss.critic_loss, input_dict)
+        if self.config["twin_q"]:
+            self.loss.critic_loss = self.twin_q_model.custom_loss(
+                self.loss.critic_loss, input_dict)
+
         TFPolicyGraph.__init__(
             self,
             observation_space,
diff --git a/python/ray/rllib/evaluation/tf_policy_graph.py b/python/ray/rllib/evaluation/tf_policy_graph.py
index 377377a60..b73395489 100644
--- a/python/ray/rllib/evaluation/tf_policy_graph.py
+++ b/python/ray/rllib/evaluation/tf_policy_graph.py
@@ -108,12 +108,6 @@ class TFPolicyGraph(PolicyGraph):
         self._prev_action_input = prev_action_input
         self._prev_reward_input = prev_reward_input
         self._sampler = action_sampler
-        if self.model:
-            self._loss = self.model.custom_loss(loss)
-            self._stats_fetches = {"model": self.model.custom_stats()}
-        else:
-            self._loss = loss
-            self._stats_fetches = {}
         self._loss_inputs = loss_inputs
         self._loss_input_dict = dict(self._loss_inputs)
         self._is_training = self._get_is_training_placeholder()
@@ -126,6 +120,13 @@ class TFPolicyGraph(PolicyGraph):
         self._max_seq_len = max_seq_len
         self._batch_divisibility_req = batch_divisibility_req
 
+        if self.model:
+            self._loss = self.model.custom_loss(loss, self._loss_input_dict)
+            self._stats_fetches = {"model": self.model.custom_stats()}
+        else:
+            self._loss = loss
+            self._stats_fetches = {}
+
         self._optimizer = self.optimizer()
         self._grads_and_vars = [(g, v)
                                 for (g, v) in self.gradients(self._optimizer)
diff --git a/python/ray/rllib/examples/custom_loss.py b/python/ray/rllib/examples/custom_loss.py
index 1b2e80cb0..4c0a21034 100644
--- a/python/ray/rllib/examples/custom_loss.py
+++ b/python/ray/rllib/examples/custom_loss.py
@@ -43,7 +43,7 @@ class CustomLossModel(Model):
                                            num_outputs, options)
         return self.fcnet.outputs, self.fcnet.last_layer
 
-    def custom_loss(self, policy_loss):
+    def custom_loss(self, policy_loss, loss_inputs):
         # create a new input reader per worker
         reader = JsonReader(self.options["custom_options"]["input_files"])
         input_ops = reader.tf_input_ops()
@@ -59,7 +59,9 @@ class CustomLossModel(Model):
         # You can also add self-supervised losses easily by referencing tensors
         # created during _build_layers_v2(). For example, an autoencoder-style
         # loss can be added as follows:
-        # ae_loss = squared_diff(self.obs_in, Decoder(self.fcnet.last_layer))
+        # ae_loss = squared_diff(
+        #     loss_inputs["obs"], Decoder(self.fcnet.last_layer))
+        print("FYI: You can also use these tensors: {}, ".format(loss_inputs))
 
         # compute the IL loss
         action_dist = Categorical(logits)
diff --git a/python/ray/rllib/models/model.py b/python/ray/rllib/models/model.py
index cb2c7c6aa..39324ee81 100644
--- a/python/ray/rllib/models/model.py
+++ b/python/ray/rllib/models/model.py
@@ -146,7 +146,7 @@ class Model(object):
             linear(self.last_layer, 1, "value", normc_initializer(1.0)), [-1])
 
     @PublicAPI
-    def custom_loss(self, policy_loss):
+    def custom_loss(self, policy_loss, loss_inputs):
         """Override to customize the loss function used to optimize this model.
 
         This can be used to incorporate self-supervised losses (by defining
@@ -158,6 +158,7 @@ class Model(object):
 
         Arguments:
             policy_loss (Tensor): scalar policy loss from the policy graph.
+            loss_inputs (dict): map of input placeholders for rollout data.
 
         Returns:
             Scalar tensor for the customized loss for this model.
diff --git a/python/ray/rllib/offline/input_reader.py b/python/ray/rllib/offline/input_reader.py
index b45126609..0e325b776 100644
--- a/python/ray/rllib/offline/input_reader.py
+++ b/python/ray/rllib/offline/input_reader.py
@@ -42,7 +42,7 @@ class InputReader(object):
 
         Example:
             >>> class MyModel(rllib.model.Model):
-            ...     def custom_loss(self, policy_loss):
+            ...     def custom_loss(self, policy_loss, loss_inputs):
             ...         reader = JsonReader(...)
             ...         input_ops = reader.tf_input_ops()
             ...         with tf.variable_scope(