class_docstr: A `PolicySaver` allows you to save a `tf_policy.Policy` to `SavedModel`.
  
  The `save()` method exports a saved model to the requested export location.
  The SavedModel that is exported can be loaded via
  `tf.compat.v2.saved_model.load` (or `tf.saved_model.load` in TF2).  The
  following signatures (concrete functions) are available: `action`,
  `get_initial_state`, and `get_train_step`.
  
  The attribute `model_variables` is also available when the saved_model is
  loaded which gives access to model variables in order to update them if
  needed.
  
  Usage:
  ```python
  
  my_policy = agent.collect_policy
  saver = PolicySaver(my_policy, batch_size=None)
  
  for i in range(...):
    agent.train(...)
    if i % 100 == 0:
      saver.save('policy_%d' % global_step)
  ```
  
  To load and use the saved policy directly:
  
  ```python
  saved_policy = tf.compat.v2.saved_model.load('policy_0')
  policy_state = saved_policy.get_initial_state(batch_size=3)
  time_step = ...
  while True:
    policy_step = saved_policy.action(time_step, policy_state)
    policy_state = policy_step.state
    time_step = f(policy_step.action)
    ...
  ```
  
  or to use the distributional form, e.g.:
  
  ```python
  batch_size = 3
  saved_policy = tf.compat.v2.saved_model.load('policy_0')
  policy_state = saved_policy.get_initial_state(batch_size=batch_size)
  time_step = ...
  while True:
    policy_step = saved_policy.distribution(time_step, policy_state)
    policy_state = policy_step.state
    time_step = f(policy_step.action.sample(batch_size))
    ...
  ```
  
  If using the flattened (signature) version, you will be limited to using
  dicts keyed by the specs' name fields.
  
  ```python
  saved_policy = tf.compat.v2.saved_model.load('policy_0')
  get_initial_state_fn = saved_policy.signatures['get_initial_state']
  action_fn = saved_policy.signatures['action']
  
  policy_state_dict = get_initial_state_fn(batch_size=3)
  time_step_dict = ...
  while True:
    time_step_state = dict(time_step_dict)
    time_step_state.update(policy_state_dict)
    policy_step_dict = action_fn(time_step_state)
    policy_state_dict = extract_policy_state_fields(policy_step_dict)
    action_dict = extract_action_fields(policy_step_dict)
    time_step_dict = f(action_dict)
    ...
  ```
docstr: Saves the policy as a checkpoint to the given `export_dir`.
  
  This will only work with checkpoints generated in TF2.x.
  
  For the checkpoint to be useful users should first call `save` to generate a
  saved_model of the policy. Checkpoints can then be used to update the policy
  without having to reload the saved_model, or saving multiple copies of the
  `saved_model.pb` file.
  
  The checkpoint is always created in the sub-directory 'variables/' and the
  checkpoint file prefix used is 'variables'. The checkpoint files are as
  follows:
     * export_dir/variables/variables.index
     * export_dir/variables/variables-xxxxx-of-xxxxx
  
  This makes the files compatible with the checkpoint part of full saved
  models, which enables you to load a saved model made up from the graph part
  of a full saved model and the variables part of a checkpoint.
  
  Args:
    export_dir: Directory to save the checkpoint to.
    options: Optional `tf.train.CheckpointOptions` object.
code: def save_checkpoint(
      self,
      export_dir: Text,
      options: Optional[tf.train.CheckpointOptions] = None,
  ):
      """Saves the policy as a checkpoint to the given `export_dir`.
  
      This will only work with checkpoints generated in TF2.x.
  
      For the checkpoint to be useful users should first call `save` to generate a
      saved_model of the policy. Checkpoints can then be used to update the policy
      without having to reload the saved_model, or saving multiple copies of the
      `saved_model.pb` file.
  
      The checkpoint is always created in the sub-directory 'variables/' and the
      checkpoint file prefix used is 'variables'. The checkpoint files are as
      follows:
         * export_dir/variables/variables.index
         * export_dir/variables/variables-xxxxx-of-xxxxx
  
      This makes the files compatible with the checkpoint part of full saved
      models, which enables you to load a saved model made up from the graph part
      of a full saved model and the variables part of a checkpoint.
  
      Args:
        export_dir: Directory to save the checkpoint to.
        options: Optional `tf.train.CheckpointOptions` object.
      """
      # In addition to the policy, also list dependencies on model_variables and
      # train_step so the checkpoint can be combined with a saved graph from a
      # full saved model.
      checkpoint = tf.compat.v2.train.Checkpoint(
          policy=self._policy,
          model_variables=self._policy.model_variables,
          train_step=self._train_step,
      )
      # Use write() to make sure that the file prefix is not modified by appending
      # a save counter value.
      file_prefix = os.path.join(
          export_dir,
          tf.saved_model.VARIABLES_DIRECTORY,
          tf.saved_model.VARIABLES_FILENAME,
      )
      checkpoint.write(file_prefix, options=options)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import functools
  import os
  from typing import Any, Callable, Dict, Tuple, Optional, Text, cast, Sequence
  from absl import logging
  import tensorflow
  import tensorflow_probability
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
  from tensorflow.python.eager import def_function
label_desc: Use image node_id 12 for calling the PolicySaver obj's underlying member method code with example usage: obj.save_checkpoint(export_dir, options) without return types