[rllib] Custom supervised loss API (#4083)

python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py

@@ -98,7 +98,8 @@ class A3CPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             obs_input=self.observations,
             action_sampler=action_dist.sample(),
             action_prob=action_dist.sampled_action_prob(),
-            loss=self.model.loss() + self.loss.total_loss,
+            loss=self.loss.total_loss,
+            model=self.model,
             loss_inputs=loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,

python/ray/rllib/agents/agent.py

@@ -243,7 +243,7 @@ class Agent(Trainable):
         self.global_vars = {"timestep": 0}
 
         # Agents allow env ids to be passed directly to the constructor.
-        self._env_id = _register_if_needed(env or config.get("env"))
+        self._env_id = self._register_if_needed(env or config.get("env"))
 
         # Create a default logger creator if no logger_creator is specified
         if logger_creator is None:
@@ -671,11 +671,14 @@ class Agent(Trainable):
         if "optimizer" in state:
             self.optimizer.restore(state["optimizer"])
 
-
-def _register_if_needed(env_object):
-    if isinstance(env_object, six.string_types):
-        return env_object
-    elif isinstance(env_object, type):
-        name = env_object.__name__
-        register_env(name, lambda config: env_object(config))
-        return name
+    def _register_if_needed(self, env_object):
+        if isinstance(env_object, six.string_types):
+            return env_object
+        elif isinstance(env_object, type):
+            name = env_object.__name__
+            register_env(name, lambda config: env_object(config))
+            return name
+        raise ValueError(
+            "{} is an invalid env specification. ".format(env_object) +
+            "You can specify a custom env as either a class "
+            "(e.g., YourEnvCls) or a registered env id (e.g., \"your_env\").")

python/ray/rllib/agents/ddpg/ddpg_policy_graph.py

@@ -334,10 +334,11 @@ class DDPGPolicyGraph(TFPolicyGraph):
                             config["l2_reg"] * 0.5 * tf.nn.l2_loss(var))
 
         # Model self-supervised losses
-        self.loss.actor_loss += self.p_model.loss()
-        self.loss.critic_loss += self.q_model.loss()
+        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.loss()
+            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

python/ray/rllib/agents/dqn/dqn_policy_graph.py

@@ -410,7 +410,8 @@ class DQNPolicyGraph(TFPolicyGraph):
             obs_input=self.cur_observations,
             action_sampler=self.output_actions,
             action_prob=self.action_prob,
-            loss=model.loss() + self.loss.loss,
+            loss=self.loss.loss,
+            model=model,
             loss_inputs=self.loss_inputs,
             update_ops=q_batchnorm_update_ops)
         self.sess.run(tf.global_variables_initializer())

python/ray/rllib/agents/impala/vtrace_policy_graph.py

@@ -216,7 +216,8 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
             obs_input=observations,
             action_sampler=action_dist.sample(),
             action_prob=action_dist.sampled_action_prob(),
-            loss=self.model.loss() + self.loss.total_loss,
+            loss=self.loss.total_loss,
+            model=self.model,
             loss_inputs=loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,

python/ray/rllib/agents/marwil/marwil_policy_graph.py

@@ -114,7 +114,8 @@ class MARWILPolicyGraph(TFPolicyGraph):
             obs_input=self.obs_t,
             action_sampler=self.output_actions,
             action_prob=action_dist.sampled_action_prob(),
-            loss=self.model.loss() + objective,
+            loss=objective,
+            model=self.model,
             loss_inputs=self.loss_inputs,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,

python/ray/rllib/agents/mock.py

@@ -46,6 +46,9 @@ class _MockAgent(Agent):
         self.info = info
         self.restored = True
 
+    def _register_if_needed(self, env_object):
+        pass
+
     def set_info(self, info):
         self.info = info
         return info

python/ray/rllib/agents/pg/pg_policy_graph.py

@@ -68,8 +68,9 @@ class PGPolicyGraph(TFPolicyGraph):
             obs_input=obs,
             action_sampler=action_dist.sample(),
             action_prob=action_dist.sampled_action_prob(),
-            loss=self.model.loss() + loss,
+            loss=loss,
             loss_inputs=loss_in,
+            model=self.model,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
             prev_action_input=prev_actions,

python/ray/rllib/agents/ppo/appo_policy_graph.py

@@ -321,7 +321,8 @@ class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             obs_input=observations,
             action_sampler=action_dist.sample(),
             action_prob=action_dist.sampled_action_prob(),
-            loss=self.model.loss() + self.loss.total_loss,
+            loss=self.loss.total_loss,
+            model=self.model,
             loss_inputs=loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
@@ -339,7 +340,6 @@ class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             values_batched = to_batches(values)
         self.stats_fetches = {
             "stats": {
-                "model_loss": self.model.loss(),
                 "cur_lr": tf.cast(self.cur_lr, tf.float64),
                 "policy_loss": self.loss.pi_loss,
                 "entropy": self.loss.entropy,

python/ray/rllib/agents/ppo/ppo_policy_graph.py

@@ -148,6 +148,8 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             existing_state_in = None
             existing_seq_lens = None
         self.observations = obs_ph
+        self.prev_actions = prev_actions_ph
+        self.prev_rewards = prev_rewards_ph
 
         self.loss_in = [
             ("obs", obs_ph),
@@ -245,7 +247,8 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             obs_input=obs_ph,
             action_sampler=self.sampler,
             action_prob=curr_action_dist.sampled_action_prob(),
-            loss=self.model.loss() + self.loss_obj.loss,
+            loss=self.loss_obj.loss,
+            model=self.model,
             loss_inputs=self.loss_in,
             state_inputs=self.model.state_in,
             state_outputs=self.model.state_out,
@@ -289,7 +292,9 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
             next_state = []
             for i in range(len(self.model.state_in)):
                 next_state.append([sample_batch["state_out_{}".format(i)][-1]])
-            last_r = self._value(sample_batch["new_obs"][-1], *next_state)
+            last_r = self._value(sample_batch["new_obs"][-1],
+                                 sample_batch["actions"][-1],
+                                 sample_batch["rewards"][-1], *next_state)
         batch = compute_advantages(
             sample_batch,
             last_r,
@@ -336,8 +341,13 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
         self.kl_coeff.load(self.kl_coeff_val, session=self.sess)
         return self.kl_coeff_val
 
-    def _value(self, ob, *args):
-        feed_dict = {self.observations: [ob], self.model.seq_lens: [1]}
+    def _value(self, ob, prev_action, prev_reward, *args):
+        feed_dict = {
+            self.observations: [ob],
+            self.prev_actions: [prev_action],
+            self.prev_rewards: [prev_reward],
+            self.model.seq_lens: [1]
+        }
         assert len(args) == len(self.model.state_in), \
             (args, self.model.state_in)
         for k, v in zip(self.model.state_in, args):