class_docstr: Triggers saves policy checkpoints an agent's policy.
  
  On construction this trigger will generate a saved_model for a:
  `greedy_policy`, a `collect_policy`, and a `raw_policy`. When triggered a
  checkpoint will be saved which can be used to updated any of the saved_model
  policies.
class_docstr_tok: ["Triggers","saves","policy","checkpoints","an","agent","'s","policy",".","On","construction","this","trigger","will","generate","a","saved_model","for","a",":","`","greedy_policy","`",",","a","`","collect_policy","`",",","and","a","`","raw_policy","`",".","When","triggered","a","checkpoint","will","be","saved","which","can","be","used","to","updated","any","of","the","saved_model","policies","."]
docstr: Initializes a PolicySavedModelTrigger.
  
  Args:
    saved_model_dir: Base dir where checkpoints will be saved.
    agent: Agent to extract policies from.
    train_step: `tf.Variable` which keeps track of the number of train steps.
    interval: How often, in train_steps, the trigger will save. Note that as
      long as the >= `interval` number of steps have passed since the last
      trigger, the event gets triggered. The current value is not necessarily
      `interval` steps away from the last triggered value.
    async_saving: If True saving will be done asynchronously in a separate
      thread. Note if this is on the variable values in the saved
      checkpoints/models are not deterministic.
    metadata_metrics: A dictionary of metrics, whose `result()` method returns
      a scalar to be saved along with the policy. Currently only supported
      when async_saving is False.
    start: Initial value for the trigger passed directly to the base class. It
      helps control from which train step the weigts of the model are saved.
    extra_concrete_functions: Optional sequence of extra concrete functions to
      register in the policy savers. The sequence should consist of tuples
      with string name for the function and the tf.function to register. Note
      this does not support adding extra assets.
    batch_size: The number of batch entries the policy will process at a time.
      This must be either `None` (unknown batch size) or a python integer.
    use_nest_path_signatures: SavedModel spec signatures will be created based
      on the sructure of the specs. Otherwise all specs must have unique
      names.
    save_greedy_policy: Disable when an agent's policy distribution method
      does not support mode.
    save_collect_policy: Disable when not saving collect policy.
    input_fn_and_spec: A `(input_fn, tensor_spec)` tuple where input_fn is a
      function that takes inputs according to tensor_spec and converts them to
      the `(time_step, policy_state)` tuple that is used as the input to the
      action_fn. When `input_fn_and_spec` is set, `tensor_spec` is the input
      for the action signature. When `input_fn_and_spec is None`, the action
      signature takes as input `(time_step, policy_state)`.
docstr_tok: ["Initializes","a","PolicySavedModelTrigger",".","Args",":","saved_model_dir",":","Base","dir","where","checkpoints","will","be","saved",".","agent",":","Agent","to","extract","policies","from",".","train_step",":","`","tf.Variable","`","which","keeps","track","of","the","number","of","train","steps",".","interval",":","How","often",",","in","train_steps",",","the","trigger","will","save",".","Note","that","as","long","as","the",">","=","`","interval","`","number","of","steps","have","passed","since","the","last","trigger",",","the","event","gets","triggered",".","The","current","value","is","not","necessarily","`","interval","`","steps","away","from","the","last","triggered","value",".","async_saving",":","If","True","saving","will","be","done","asynchronously","in","a","separate","thread",".","Note","if","this","is","on","the","variable","values","in","the","saved","checkpoints\/models","are","not","deterministic",".","metadata_metrics",":","A","dictionary","of","metrics",",","whose","`","result","(",")","`","method","returns","a","scalar","to","be","saved","along","with","the","policy",".","Currently","only","supported","when","async_saving","is","False",".","start",":","Initial","value","for","the","trigger","passed","directly","to","the","base","class",".","It","helps","control","from","which","train","step","the","weigts","of","the","model","are","saved",".","extra_concrete_functions",":","Optional","sequence","of","extra","concrete","functions","to","register","in","the","policy","savers",".","The","sequence","should","consist","of","tuples","with","string","name","for","the","function","and","the","tf.function","to","register",".","Note","this","does","not","support","adding","extra","assets",".","batch_size",":","The","number","of","batch","entries","the","policy","will","process","at","a","time",".","This","must","be","either","`","None","`","(","unknown","batch","size",")","or","a","python","integer",".","use_nest_path_signatures",":","SavedModel","spec","signatures","will","be","created","based","on","the","sructure","of","the","specs",".","Otherwise","all","specs","must","have","unique","names",".","save_greedy_policy",":","Disable","when","an","agent","'s","policy","distribution","method","does","not","support","mode",".","save_collect_policy",":","Disable","when","not","saving","collect","policy",".","input_fn_and_spec",":","A","`","(","input_fn",",","tensor_spec",")","`","tuple","where","input_fn","is","a","function","that","takes","inputs","according","to","tensor_spec","and","converts","them","to","the","`","(","time_step",",","policy_state",")","`","tuple","that","is","used","as","the","input","to","the","action_fn",".","When","`","input_fn_and_spec","`","is","set",",","`","tensor_spec","`","is","the","input","for","the","action","signature",".","When","`","input_fn_and_spec","is","None","`",",","the","action","signature","takes","as","input","`","(","time_step",",","policy_state",")","`","."]
code: def __init__(
      self,
      saved_model_dir: Text,
      agent: tf_agent.TFAgent,
      train_step: tf.Variable,
      interval: int,
      async_saving: bool = False,
      metadata_metrics: Optional[
          Mapping[Text, py_metric.PyMetric]
      ] = None,
      start: int = 0,
      extra_concrete_functions: Optional[
          Sequence[Tuple[str, policy_saver.def_function.Function]]
      ] = None,
      batch_size: Optional[int] = None,
      use_nest_path_signatures: bool = True,
      save_greedy_policy=True,
      save_collect_policy=True,
      input_fn_and_spec: Optional[
          policy_saver.InputFnAndSpecType
      ] = None,
  ):
      """Initializes a PolicySavedModelTrigger.
  
      Args:
        saved_model_dir: Base dir where checkpoints will be saved.
        agent: Agent to extract policies from.
        train_step: `tf.Variable` which keeps track of the number of train steps.
        interval: How often, in train_steps, the trigger will save. Note that as
          long as the >= `interval` number of steps have passed since the last
          trigger, the event gets triggered. The current value is not necessarily
          `interval` steps away from the last triggered value.
        async_saving: If True saving will be done asynchronously in a separate
          thread. Note if this is on the variable values in the saved
          checkpoints/models are not deterministic.
        metadata_metrics: A dictionary of metrics, whose `result()` method returns
          a scalar to be saved along with the policy. Currently only supported
          when async_saving is False.
        start: Initial value for the trigger passed directly to the base class. It
          helps control from which train step the weigts of the model are saved.
        extra_concrete_functions: Optional sequence of extra concrete functions to
          register in the policy savers. The sequence should consist of tuples
          with string name for the function and the tf.function to register. Note
          this does not support adding extra assets.
        batch_size: The number of batch entries the policy will process at a time.
          This must be either `None` (unknown batch size) or a python integer.
        use_nest_path_signatures: SavedModel spec signatures will be created based
          on the sructure of the specs. Otherwise all specs must have unique
          names.
        save_greedy_policy: Disable when an agent's policy distribution method
          does not support mode.
        save_collect_policy: Disable when not saving collect policy.
        input_fn_and_spec: A `(input_fn, tensor_spec)` tuple where input_fn is a
          function that takes inputs according to tensor_spec and converts them to
          the `(time_step, policy_state)` tuple that is used as the input to the
          action_fn. When `input_fn_and_spec` is set, `tensor_spec` is the input
          for the action signature. When `input_fn_and_spec is None`, the action
          signature takes as input `(time_step, policy_state)`.
      """
      if async_saving and metadata_metrics:
          raise NotImplementedError(
              "Support for metadata_metrics is not "
              "implemented for async policy saver."
          )
  
      self._agent = agent
      self._train_step = train_step
      self._async_saving = async_saving
      self._metadata_metrics = metadata_metrics or {}
      self._metadata = {
          k: tf.Variable(
              0, dtype=v.result().dtype, shape=v.result().shape
          )
          for k, v in self._metadata_metrics.items()
      }
      self._input_fn_and_spec = input_fn_and_spec
  
      greedy = None
      if isinstance(agent.policy, greedy_policy.GreedyPolicy):
          raw_policy = agent.policy.wrapped_policy
          greedy = agent.policy
      else:
          raw_policy = agent.policy
          if save_greedy_policy:
              greedy = greedy_policy.GreedyPolicy(agent.policy)
  
      self._raw_policy_saver = self._build_saver(
          raw_policy, batch_size, use_nest_path_signatures
      )
      savers = [
          (self._raw_policy_saver, learner.RAW_POLICY_SAVED_MODEL_DIR)
      ]
  
      if save_collect_policy:
          collect_policy_saver = self._build_saver(
              agent.collect_policy, batch_size, use_nest_path_signatures
          )
  
          savers.append(
              (
                  collect_policy_saver,
                  learner.COLLECT_POLICY_SAVED_MODEL_DIR,
              )
          )
  
      if save_greedy_policy:
          greedy_policy_saver = self._build_saver(
              greedy, batch_size, use_nest_path_signatures
          )
          savers.append(
              (
                  greedy_policy_saver,
                  learner.GREEDY_POLICY_SAVED_MODEL_DIR,
              )
          )
  
      extra_concrete_functions = extra_concrete_functions or []
      for saver, _ in savers:
          for name, fn in extra_concrete_functions:
              saver.register_concrete_function(name, fn)
  
      self._checkpoint_dir = os.path.join(
          saved_model_dir, learner.POLICY_CHECKPOINT_DIR
      )
  
      # TODO(b/173815037): Use a TF-Agents util to check for whether a saved
      # policy already exists.
      for saver, path in savers:
          spec_path = os.path.join(
              saved_model_dir, path, "policy_specs.pbtxt"
          )
          if not tf.io.gfile.exists(spec_path):
              saver.save(os.path.join(saved_model_dir, path))
  
      super(PolicySavedModelTrigger, self).__init__(
          interval, self._save_fn, start=start
      )
  
imports: import os
  from typing import Mapping, Optional, Sequence, Text, Tuple, Union
  from absl import logging
  import tensorflow.compat.v2
  from tf_agents.agents import tf_agent
  from tf_agents.metrics import py_metric
  from tf_agents.policies import async_policy_saver
  from tf_agents.policies import greedy_policy
  from tf_agents.policies import policy_saver
  from tf_agents.policies import tf_policy
  from tf_agents.train import interval_trigger
  from tf_agents.train import learner
  from tf_agents.train import step_per_second_tracker
  from tf_agents.typing import types
label_desc: Use image node_id 1 to create a new PolicySavedModelTrigger object from inherited base classes: interval_trigger with example: obj = PolicySavedModelTrigger(saved_model_dir, agent, train_step, interval, async_saving, metadata_metrics, start, extra_concrete_functions, batch_size, use_nest_path_signatures, save_greedy_policy, save_collect_policy, input_fn_and_spec)