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: Spec that describes the output of `get_initial_state` in the SavedModel.
  
  This may differ from the original policy if `use_nest_path_signatures` was
  enabled.
  
  Returns:
    A nest of specs.
code: def policy_state_spec(self) -> types.NestedTensorSpec:
      """Spec that describes the output of `get_initial_state` in the SavedModel.
  
      This may differ from the original policy if `use_nest_path_signatures` was
      enabled.
  
      Returns:
        A nest of specs.
      """
      return self._policy_state_spec
  
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 5 for calling the PolicySaver obj's underlying member method code with example usage: obj.policy_state_spec() and returns: self