From d90f3653946c0be3edaad68fbb2660d164099bf1 Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Mon, 12 Nov 2018 18:55:24 -0800
Subject: [PATCH] [rllib] Add self-supervised loss to model (#3291)

# What do these changes do?

Allow self-supervised losses to be easily defined in custom models. Add this to the reference policy graphs.
---
 doc/source/rllib-models.rst                   | 14 ++++++++++++-
 .../rllib/agents/a3c/a3c_tf_policy_graph.py   |  2 +-
 .../rllib/agents/ddpg/ddpg_policy_graph.py    | 20 ++++++++++---------
 .../ray/rllib/agents/dqn/dqn_policy_graph.py  | 13 ++++++------
 .../agents/impala/vtrace_policy_graph.py      |  2 +-
 python/ray/rllib/agents/pg/pg_policy_graph.py |  2 +-
 .../ray/rllib/agents/ppo/ppo_policy_graph.py  |  2 +-
 python/ray/rllib/models/model.py              | 12 +++++++++++
 8 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/doc/source/rllib-models.rst b/doc/source/rllib-models.rst
index 79df3ab5c..5fde37f53 100644
--- a/doc/source/rllib-models.rst
+++ b/doc/source/rllib-models.rst
@@ -30,7 +30,7 @@ The following is a list of the built-in model hyperparameters:
 Custom Models
 -------------
 
-Custom models should subclass the common RLlib `model class <https://github.com/ray-project/ray/blob/master/python/ray/rllib/models/model.py>`__ and override the ``_build_layers_v2`` method. This method takes in a dict of tensor inputs (the observation ``obs``, ``prev_action``, and ``prev_reward``), and returns a feature layer and float vector of the specified output size. You can also override the ``value_function`` method to implement a custom value branch. The model can then be registered and used in place of a built-in model:
+Custom models should subclass the common RLlib `model class <https://github.com/ray-project/ray/blob/master/python/ray/rllib/models/model.py>`__ and override the ``_build_layers_v2`` method. This method takes in a dict of tensor inputs (the observation ``obs``, ``prev_action``, and ``prev_reward``), and returns a feature layer and float vector of the specified output size. You can also override the ``value_function`` method to implement a custom value branch. A self-supervised loss can be defined via the ``loss`` method. The model can then be registered and used in place of a built-in model:
 
 .. code-block:: python
 
@@ -86,6 +86,18 @@ Custom models should subclass the common RLlib `model class <https://github.com/
             return tf.reshape(
                 linear(self.last_layer, 1, "value", normc_initializer(1.0)), [-1])
 
+        def loss(self):
+            """Builds any built-in (self-supervised) loss for the model.
+
+            For example, this can be used to incorporate auto-encoder style losses.
+            Note that this loss has to be included in the policy graph loss to have
+            an effect (done for built-in algorithms).
+
+            Returns:
+                Scalar tensor for the self-supervised loss.
+            """
+            return tf.constant(0.0)
+
     ModelCatalog.register_custom_model("my_model", MyModelClass)
 
     ray.init()
diff --git a/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py b/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py
index e04d418c6..6f079713a 100644
--- a/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py
+++ b/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py
@@ -94,7 +94,7 @@ class A3CPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             self.sess,
             obs_input=self.observations,
             action_sampler=action_dist.sample(),
-            loss=self.loss.total_loss,
+            loss=self.model.loss() + self.loss.total_loss,
             loss_inputs=loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
diff --git a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
index e647a4153..5d1ccf2e3 100644
--- a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
+++ b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
@@ -86,6 +86,7 @@ class QNetwork(object):
                 q_out, num_outputs=hidden, activation_fn=activation)
         self.value = layers.fully_connected(
             q_out, num_outputs=1, activation_fn=None)
+        self.model = model
 
 
 class ActorCriticLoss(object):
@@ -198,17 +199,17 @@ class DDPGPolicyGraph(TFPolicyGraph):
 
         # q network evaluation
         with tf.variable_scope(Q_SCOPE) as scope:
-            q_t = self._build_q_network(self.obs_t, observation_space,
-                                        self.act_t)
+            q_t, model = self._build_q_network(self.obs_t, observation_space,
+                                               self.act_t)
             self.q_func_vars = _scope_vars(scope.name)
         with tf.variable_scope(Q_SCOPE, reuse=True):
-            q_tp0 = self._build_q_network(self.obs_t, observation_space,
-                                          output_actions)
+            q_tp0, _ = self._build_q_network(self.obs_t, observation_space,
+                                             output_actions)
 
         # target q network evalution
         with tf.variable_scope(Q_TARGET_SCOPE) as scope:
-            q_tp1 = self._build_q_network(self.obs_tp1, observation_space,
-                                          output_actions_estimated)
+            q_tp1, _ = self._build_q_network(self.obs_tp1, observation_space,
+                                             output_actions_estimated)
             target_q_func_vars = _scope_vars(scope.name)
 
         self.loss = self._build_actor_critic_loss(q_t, q_tp1, q_tp0)
@@ -258,7 +259,7 @@ class DDPGPolicyGraph(TFPolicyGraph):
             self.sess,
             obs_input=self.cur_observations,
             action_sampler=self.output_actions,
-            loss=self.loss.total_loss,
+            loss=model.loss() + self.loss.total_loss,
             loss_inputs=self.loss_inputs)
         self.sess.run(tf.global_variables_initializer())
 
@@ -271,12 +272,13 @@ class DDPGPolicyGraph(TFPolicyGraph):
         self.update_target(tau=1.0)
 
     def _build_q_network(self, obs, obs_space, actions):
-        return QNetwork(
+        q_net = QNetwork(
             ModelCatalog.get_model({
                 "obs": obs
             }, obs_space, 1, self.config["model"]), actions,
             self.config["critic_hiddens"],
-            self.config["critic_hidden_activation"]).value
+            self.config["critic_hidden_activation"])
+        return q_net.value, q_net.model
 
     def _build_p_network(self, obs, obs_space):
         return PNetwork(
diff --git a/python/ray/rllib/agents/dqn/dqn_policy_graph.py b/python/ray/rllib/agents/dqn/dqn_policy_graph.py
index f279d4ecc..6125cd9d3 100644
--- a/python/ray/rllib/agents/dqn/dqn_policy_graph.py
+++ b/python/ray/rllib/agents/dqn/dqn_policy_graph.py
@@ -28,6 +28,7 @@ class QNetwork(object):
                  v_min=-10.0,
                  v_max=10.0,
                  sigma0=0.5):
+        self.model = model
         with tf.variable_scope("action_value"):
             action_out = model.last_layer
             for i in range(len(hiddens)):
@@ -274,7 +275,7 @@ class DQNPolicyGraph(TFPolicyGraph):
 
         # Action Q network
         with tf.variable_scope(Q_SCOPE) as scope:
-            q_values, q_logits, q_dist = self._build_q_network(
+            q_values, q_logits, q_dist, _ = self._build_q_network(
                 self.cur_observations, observation_space)
             self.q_func_vars = _scope_vars(scope.name)
 
@@ -294,12 +295,12 @@ class DQNPolicyGraph(TFPolicyGraph):
 
         # q network evaluation
         with tf.variable_scope(Q_SCOPE, reuse=True):
-            q_t, q_logits_t, q_dist_t = self._build_q_network(
+            q_t, q_logits_t, q_dist_t, model = self._build_q_network(
                 self.obs_t, observation_space)
 
         # target q network evalution
         with tf.variable_scope(Q_TARGET_SCOPE) as scope:
-            q_tp1, q_logits_tp1, q_dist_tp1 = self._build_q_network(
+            q_tp1, q_logits_tp1, q_dist_tp1, _ = self._build_q_network(
                 self.obs_tp1, observation_space)
             self.target_q_func_vars = _scope_vars(scope.name)
 
@@ -313,7 +314,7 @@ class DQNPolicyGraph(TFPolicyGraph):
         if config["double_q"]:
             with tf.variable_scope(Q_SCOPE, reuse=True):
                 q_tp1_using_online_net, q_logits_tp1_using_online_net, \
-                    q_dist_tp1_using_online_net = self._build_q_network(
+                    q_dist_tp1_using_online_net, _ = self._build_q_network(
                         self.obs_tp1, observation_space)
             q_tp1_best_using_online_net = tf.argmax(q_tp1_using_online_net, 1)
             q_tp1_best_one_hot_selection = tf.one_hot(
@@ -359,7 +360,7 @@ class DQNPolicyGraph(TFPolicyGraph):
             self.sess,
             obs_input=self.cur_observations,
             action_sampler=self.output_actions,
-            loss=self.loss.loss,
+            loss=model.loss() + self.loss.loss,
             loss_inputs=self.loss_inputs)
         self.sess.run(tf.global_variables_initializer())
 
@@ -371,7 +372,7 @@ class DQNPolicyGraph(TFPolicyGraph):
             self.config["dueling"], self.config["hiddens"],
             self.config["noisy"], self.config["num_atoms"],
             self.config["v_min"], self.config["v_max"], self.config["sigma0"])
-        return qnet.value, qnet.logits, qnet.dist
+        return qnet.value, qnet.logits, qnet.dist, qnet.model
 
     def _build_q_value_policy(self, q_values):
         return QValuePolicy(q_values, self.cur_observations, self.num_actions,
diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py
index b75d79e57..a08836e93 100644
--- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py
+++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py
@@ -203,7 +203,7 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
             self.sess,
             obs_input=observations,
             action_sampler=action_dist.sample(),
-            loss=self.loss.total_loss,
+            loss=self.model.loss() + self.loss.total_loss,
             loss_inputs=loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
diff --git a/python/ray/rllib/agents/pg/pg_policy_graph.py b/python/ray/rllib/agents/pg/pg_policy_graph.py
index 0395f026e..178cf29e6 100644
--- a/python/ray/rllib/agents/pg/pg_policy_graph.py
+++ b/python/ray/rllib/agents/pg/pg_policy_graph.py
@@ -64,7 +64,7 @@ class PGPolicyGraph(TFPolicyGraph):
             sess,
             obs_input=obs,
             action_sampler=action_dist.sample(),
-            loss=loss,
+            loss=self.model.loss() + loss,
             loss_inputs=loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
diff --git a/python/ray/rllib/agents/ppo/ppo_policy_graph.py b/python/ray/rllib/agents/ppo/ppo_policy_graph.py
index 1a78625b2..f43a33625 100644
--- a/python/ray/rllib/agents/ppo/ppo_policy_graph.py
+++ b/python/ray/rllib/agents/ppo/ppo_policy_graph.py
@@ -230,7 +230,7 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             self.sess,
             obs_input=obs_ph,
             action_sampler=self.sampler,
-            loss=self.loss_obj.loss,
+            loss=self.model.loss() + self.loss_obj.loss,
             loss_inputs=self.loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
diff --git a/python/ray/rllib/models/model.py b/python/ray/rllib/models/model.py
index 091b17c8f..0df1b5d09 100644
--- a/python/ray/rllib/models/model.py
+++ b/python/ray/rllib/models/model.py
@@ -144,6 +144,18 @@ class Model(object):
         return tf.reshape(
             linear(self.last_layer, 1, "value", normc_initializer(1.0)), [-1])
 
+    def loss(self):
+        """Builds any built-in (self-supervised) loss for the model.
+
+        For example, this can be used to incorporate auto-encoder style losses.
+        Note that this loss has to be included in the policy graph loss to have
+        an effect (done for built-in algorithms).
+
+        Returns:
+            Scalar tensor for the self-supervised loss.
+        """
+        return tf.constant(0.0)
+
 
 def _restore_original_dimensions(input_dict, obs_space):
     if hasattr(obs_space, "original_space"):