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: Save the policy to the given `export_dir`.
  
  Args:
    export_dir: Directory to save the policy to.
    options: Optional `tf.saved_model.SaveOptions` object.
code: def save(
      self,
      export_dir: Text,
      options: Optional[tf.saved_model.SaveOptions] = None,
  ):
      """Save the policy to the given `export_dir`.
  
      Args:
        export_dir: Directory to save the policy to.
        options: Optional `tf.saved_model.SaveOptions` object.
      """
      tf.compat.v2.saved_model.save(
          self._policy,
          export_dir,
          signatures=self._signatures,
          options=options,
      )
  
      temp_spec_file_name = "{}_temp".format(POLICY_SPECS_PBTXT)
      temp_spec_output_path = os.path.join(
          export_dir, temp_spec_file_name
      )
      specs = {
          "collect_data_spec": self._policy.collect_data_spec,
          "policy_state_spec": self._policy.policy_state_spec,
      }
      tensor_spec.to_pbtxt_file(temp_spec_output_path, specs)
      spec_output_path = os.path.join(export_dir, POLICY_SPECS_PBTXT)
      # By moving the file to its final location makes it safer to wait for the
      # file (e.g. from a separate binary). The parameter `overwrite=True`
      # reproduces the exact previous behavior.
      tf.io.gfile.rename(
          temp_spec_output_path, spec_output_path, overwrite=True
      )
  
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 11 for calling the PolicySaver obj's underlying member method code with example usage: obj.save(export_dir, options) without return types