code: def create_staleness_metrics_after_train_step_fn(
      train_step: tf.Variable, train_steps_per_policy_update: int = 1
  ) -> Callable[
      [
          Tuple[types.NestedTensor, types.ReverbSampleInfo],
          tf_agent.LossInfo,
      ],
      None,
  ]:
      """Creates an `after_train_step_fn` that computes staleness summaries.
  
      Staleness, in this context, means that the observation was generated by a
      policy that is older than the recently outputed policy.
      Assume that observation train step is stored as Reverb priorities.
  
      Args:
        train_step: The current train step.
        train_steps_per_policy_update: Number of train iterations to perform between
          two policy updates.
  
      Returns:
        The created `after_train_step_fn`.
      """
  
      def after_train_step_fn(experience, loss_info):
          del loss_info  # Unused.
          _, sample_info = experience
  
          # Get the train step in which the experience was observed. This is stored as
          # Reverb priority.
          # TODO(b/168426331): Check sample info version.
          observation_generation_train_step = tf.cast(
              sample_info.priority, dtype=tf.int64
          )
  
          # Get the train step corresponding to the latest outputed policy.
          # Policy is written in every `train_steps_per_policy_update` step, so we
          # normalize the value of `train_step` accordingly.
          on_policy_train_step = (
              tf.cast(
                  train_step / train_steps_per_policy_update,
                  dtype=tf.int64,
              )
              * train_steps_per_policy_update
          )
  
          # An observation is off-policy if its train step delta is greater than 0.
          observation_train_step_delta = (
              on_policy_train_step - observation_generation_train_step
          )
          max_train_step_delta = tf.reduce_max(
              observation_train_step_delta
          )
          max_policy_update_delta = tf.cast(
              max_train_step_delta / train_steps_per_policy_update,
              dtype=tf.int64,
          )
          num_stale_observations = tf.reduce_sum(
              tf.cast(observation_train_step_delta > 0, tf.int64)
          )
  
          # Break out from local name scopes (e.g. the ones intrdouced by while loop).
          with tf.name_scope(""):
              # Write the summaries for the first replica.
              tf.summary.scalar(
                  name="staleness/max_train_step_delta_in_batch",
                  data=max_train_step_delta,
                  step=train_step,
              )
              tf.summary.scalar(
                  name="staleness/max_policy_update_delta_in_batch",
                  data=max_policy_update_delta,
                  step=train_step,
              )
              tf.summary.scalar(
                  name="staleness/num_stale_obserations_in_batch",
                  data=num_stale_observations,
                  step=train_step,
              )
  
      return after_train_step_fn
  
file_path: agents/tf_agents/train/utils/train_utils.py
imports: import os
  import time
  from typing import Callable, Text, Tuple
  from absl import logging
  import tensorflow.compat.v2
  from tf_agents.agents import tf_agent
  from tf_agents.policies import py_tf_eager_policy
  from tf_agents.typing import types
  from tf_agents.utils import lazy_loader
label_desc: Use image node_id 2 for calling a global function with example usage: create_staleness_metrics_after_train_step_fn(train_step, train_steps_per_policy_update) and returns: after_train_step_fn