[RLlib] Add on_learn_on_batch (Policy) callback to DefaultCallbacks. (#12070)

rllib/agents/callbacks.py

@@ -32,7 +32,8 @@ class DefaultCallbacks:
 
     def on_episode_start(self, *, worker: "RolloutWorker", base_env: BaseEnv,
                          policies: Dict[PolicyID, Policy],
-                         episode: MultiAgentEpisode, env_index: int, **kwargs):
+                         episode: MultiAgentEpisode, env_index: int,
+                         **kwargs) -> None:
         """Callback run on the rollout worker before each episode starts.
 
         Args:
@@ -58,7 +59,8 @@ class DefaultCallbacks:
             })
 
     def on_episode_step(self, *, worker: "RolloutWorker", base_env: BaseEnv,
-                        episode: MultiAgentEpisode, env_index: int, **kwargs):
+                        episode: MultiAgentEpisode, env_index: int,
+                        **kwargs) -> None:
         """Runs on each episode step.
 
         Args:
@@ -82,7 +84,8 @@ class DefaultCallbacks:
 
     def on_episode_end(self, *, worker: "RolloutWorker", base_env: BaseEnv,
                        policies: Dict[PolicyID, Policy],
-                       episode: MultiAgentEpisode, env_index: int, **kwargs):
+                       episode: MultiAgentEpisode, env_index: int,
+                       **kwargs) -> None:
         """Runs when an episode is done.
 
         Args:
@@ -111,7 +114,7 @@ class DefaultCallbacks:
             self, *, worker: "RolloutWorker", episode: MultiAgentEpisode,
             agent_id: AgentID, policy_id: PolicyID,
             policies: Dict[PolicyID, Policy], postprocessed_batch: SampleBatch,
-            original_batches: Dict[AgentID, SampleBatch], **kwargs):
+            original_batches: Dict[AgentID, SampleBatch], **kwargs) -> None:
         """Called immediately after a policy's postprocess_fn is called.
 
         You can use this callback to do additional postprocessing for a policy,
@@ -143,7 +146,7 @@ class DefaultCallbacks:
             })
 
     def on_sample_end(self, *, worker: "RolloutWorker", samples: SampleBatch,
-                      **kwargs):
+                      **kwargs) -> None:
         """Called at the end of RolloutWorker.sample().
 
         Args:
@@ -159,7 +162,23 @@ class DefaultCallbacks:
                 "samples": samples,
             })
 
-    def on_train_result(self, *, trainer, result: dict, **kwargs):
+    def on_learn_on_batch(self, *, policy: Policy, train_batch: SampleBatch,
+                          **kwargs) -> None:
+        """Called at the beginning of Policy.learn_on_batch().
+
+        Note: This is called before the Model's `preprocess_train_batch()`
+        is called.
+
+        Args:
+            policy (Policy): Reference to the current Policy object.
+            train_batch (SampleBatch): SampleBatch to be trained on. You can
+                mutate this object to modify the samples generated.
+            kwargs: Forward compatibility placeholder.
+        """
+
+        pass
+
+    def on_train_result(self, *, trainer, result: dict, **kwargs) -> None:
         """Called at the end of Trainable.train().
 
         Args:

rllib/policy/eager_tf_policy.py

@@ -302,14 +302,18 @@ def build_eager_tf_policy(name,
             return sample_batch
 
         @override(Policy)
-        def learn_on_batch(self, samples):
+        def learn_on_batch(self, postprocessed_batch):
+            # Callback handling.
+            self.callbacks.on_learn_on_batch(
+                policy=self, train_batch=postprocessed_batch)
+
             # Get batch ready for RNNs, if applicable.
             pad_batch_to_sequences_of_same_size(
-                samples,
+                postprocessed_batch,
                 shuffle=False,
                 max_seq_len=self._max_seq_len,
                 batch_divisibility_req=self.batch_divisibility_req)
-            return self._learn_on_batch_eager(samples)
+            return self._learn_on_batch_eager(postprocessed_batch)
 
         @convert_eager_inputs
         @convert_eager_outputs

rllib/policy/policy.py

@@ -71,6 +71,11 @@ class Policy(metaclass=ABCMeta):
         self.action_space = action_space
         self.action_space_struct = get_base_struct_from_space(action_space)
         self.config = config
+        if self.config.get("callbacks"):
+            self.callbacks: "DefaultCallbacks" = self.config.get("callbacks")()
+        else:
+            from ray.rllib.agents.callbacks import DefaultCallbacks
+            self.callbacks: "DefaultCallbacks" = DefaultCallbacks()
         # The global timestep, broadcast down from time to time from the
         # driver.
         self.global_timestep = 0

rllib/policy/tf_policy.py

@@ -765,6 +765,11 @@ class TFPolicy(Policy):
 
     def _build_learn_on_batch(self, builder, postprocessed_batch):
         self._debug_vars()
+
+        # Callback handling.
+        self.callbacks.on_learn_on_batch(
+            policy=self, train_batch=postprocessed_batch)
+
         builder.add_feed_dict(self.extra_compute_grad_feed_dict())
         builder.add_feed_dict(
             self._get_loss_inputs_dict(postprocessed_batch, shuffle=False))

rllib/policy/torch_policy.py

@@ -325,6 +325,10 @@ class TorchPolicy(Policy):
     @DeveloperAPI
     def learn_on_batch(
             self, postprocessed_batch: SampleBatch) -> Dict[str, TensorType]:
+        # Callback handling.
+        self.callbacks.on_learn_on_batch(
+            policy=self, train_batch=postprocessed_batch)
+
         # Compute gradients (will calculate all losses and `backward()`
         # them to get the grads).
         grads, fetches = self.compute_gradients(postprocessed_batch)