code: def testSaveGetInitialState(self):
      network = q_rnn_network.QRnnNetwork(
          input_tensor_spec=self._time_step_spec.observation,
          action_spec=self._action_spec,
          lstm_size=(40,),
      )
  
      policy = q_policy.QPolicy(
          time_step_spec=self._time_step_spec,
          action_spec=self._action_spec,
          q_network=network,
      )
  
      train_step = common.create_variable("train_step", initial_value=0)
      saver_nobatch = policy_saver.PolicySaver(
          policy,
          train_step=train_step,
          batch_size=None,
          use_nest_path_signatures=False,
      )
      path = os.path.join(
          self.get_temp_dir(), "save_model_initial_state_nobatch"
      )
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
  
      with self.cached_session():
          saver_nobatch.save(path)
          reloaded_nobatch = tf.compat.v2.saved_model.load(path)
          self.evaluate(
              tf.compat.v1.initializers.variables(
                  reloaded_nobatch.model_variables
              )
          )
  
      self.assertIn("get_initial_state", reloaded_nobatch.signatures)
      reloaded_get_initial_state = reloaded_nobatch.signatures[
          "get_initial_state"
      ]
      self._compare_input_output_specs(
          reloaded_get_initial_state,
          expected_input_specs=(
              tf.TensorSpec(
                  dtype=tf.int32, shape=(), name="batch_size"
              ),
          ),
          expected_output_spec=policy.policy_state_spec,
          batch_input=False,
          batch_size=None,
      )
  
      initial_state = policy.get_initial_state(batch_size=3)
      initial_state = self.evaluate(initial_state)
  
      reloaded_nobatch_initial_state = (
          reloaded_nobatch.get_initial_state(batch_size=3)
      )
      reloaded_nobatch_initial_state = self.evaluate(
          reloaded_nobatch_initial_state
      )
      tf.nest.map_structure(
          self.assertAllClose,
          initial_state,
          reloaded_nobatch_initial_state,
      )
  
      saver_batch = policy_saver.PolicySaver(
          policy,
          train_step=train_step,
          batch_size=3,
          use_nest_path_signatures=False,
      )
      path = os.path.join(
          self.get_temp_dir(), "save_model_initial_state_batch"
      )
      with self.cached_session():
          saver_batch.save(path)
          reloaded_batch = tf.compat.v2.saved_model.load(path)
          self.evaluate(
              tf.compat.v1.initializers.variables(
                  reloaded_batch.model_variables
              )
          )
      self.assertIn("get_initial_state", reloaded_batch.signatures)
      reloaded_get_initial_state = reloaded_batch.signatures[
          "get_initial_state"
      ]
      self._compare_input_output_specs(
          reloaded_get_initial_state,
          expected_input_specs=(),
          expected_output_spec=policy.policy_state_spec,
          batch_input=False,
          batch_size=3,
      )
  
      reloaded_batch_initial_state = reloaded_batch.get_initial_state()
      reloaded_batch_initial_state = self.evaluate(
          reloaded_batch_initial_state
      )
      tf.nest.map_structure(
          self.assertAllClose,
          initial_state,
          reloaded_batch_initial_state,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import shutil
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.networks import actor_distribution_network
  from tf_agents.networks import q_network
  from tf_agents.networks import q_rnn_network
  from tf_agents.policies import actor_policy
  from tf_agents.policies import policy_saver
  from tf_agents.policies import q_policy
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 6 for calling the PolicySaverTest obj's underlying member method code with example usage: obj.testSaveGetInitialState() without return types