diff --git a/rllib/agents/callbacks.py b/rllib/agents/callbacks.py index 4621adef2..f0460c17d 100644 --- a/rllib/agents/callbacks.py +++ b/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: diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 39b3ef7be..7d92f6500 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/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 diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index b8cb695ae..80a2fae1b 100644 --- a/rllib/policy/policy.py +++ b/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 diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py index 4a32dca06..d5e2a8813 100644 --- a/rllib/policy/tf_policy.py +++ b/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)) diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 0c15151f7..aaae43246 100644 --- a/rllib/policy/torch_policy.py +++ b/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)