[rllib] on_train_result results do not get logged (#3865)

python/ray/rllib/agents/agent.py

@@ -292,13 +292,18 @@ class Agent(Trainable):
             logger.debug("synchronized filters: {}".format(
                 self.local_evaluator.filters))
 
+        return result
+
+    @override(Trainable)
+    def _log_result(self, result):
         if self.config["callbacks"].get("on_train_result"):
             self.config["callbacks"]["on_train_result"]({
                 "agent": self,
                 "result": result,
             })
-
-        return result
+        # log after the callback is invoked, so that the user has a chance
+        # to mutate the result
+        Trainable._log_result(self, result)
 
     @override(Trainable)
     def _setup(self, config):

python/ray/tune/trainable.py

@@ -202,7 +202,7 @@ class Trainable(object):
             timesteps_since_restore=self._timesteps_since_restore,
             iterations_since_restore=self._iterations_since_restore)
 
-        self._result_logger.on_result(result)
+        self._log_result(result)
 
         return result
 
@@ -415,6 +415,15 @@ class Trainable(object):
         """
         pass
 
+    def _log_result(self, result):
+        """Subclasses can optionally override this to customize logging.
+
+        Args:
+            result (dict): Training result returned by _train().
+        """
+
+        self._result_logger.on_result(result)
+
     def _stop(self):
         """Subclasses should override this for any cleanup on stop."""
         pass