code: def testSaveAction(
      self, seeded, has_state, distribution_net, has_input_fn_and_spec
  ):
      with tf.compat.v1.Graph().as_default():
          tf.compat.v1.set_random_seed(self._global_seed)
          with tf.compat.v1.Session().as_default():
              global_step = common.create_variable(
                  "train_step", initial_value=0
              )
              if distribution_net:
                  network = actor_distribution_network.ActorDistributionNetwork(
                      self._time_step_spec.observation,
                      self._action_spec,
                  )
                  policy = actor_policy.ActorPolicy(
                      time_step_spec=self._time_step_spec,
                      action_spec=self._action_spec,
                      actor_network=network,
                  )
              else:
                  if has_state:
                      network = q_rnn_network.QRnnNetwork(
                          input_tensor_spec=self._time_step_spec.observation,
                          action_spec=self._action_spec,
                          lstm_size=(40,),
                      )
                  else:
                      network = q_network.QNetwork(
                          input_tensor_spec=self._time_step_spec.observation,
                          action_spec=self._action_spec,
                      )
  
                  policy = q_policy.QPolicy(
                      time_step_spec=self._time_step_spec,
                      action_spec=self._action_spec,
                      q_network=network,
                  )
  
              action_seed = 98723
  
              batch_size = 3
              action_inputs = tensor_spec.sample_spec_nest(
                  (self._time_step_spec, policy.policy_state_spec),
                  outer_dims=(batch_size,),
                  seed=4,
              )
              action_input_values = self.evaluate(action_inputs)
              action_input_tensors = tf.nest.map_structure(
                  tf.convert_to_tensor, action_input_values
              )
  
              action_output = policy.action(
                  *action_input_tensors, seed=action_seed
              )
              distribution_output = policy.distribution(
                  *action_input_tensors
              )
              self.assertIsInstance(
                  distribution_output.action,
                  tfp.distributions.Distribution,
              )
  
              self.evaluate(tf.compat.v1.global_variables_initializer())
  
              action_output_dict = collections.OrderedDict(
                  (
                      (spec.name, value)
                      for (spec, value) in zip(
                          tf.nest.flatten(policy.policy_step_spec),
                          tf.nest.flatten(action_output),
                      )
                  )
              )
  
              # Check output of the flattened signature call.
              (action_output_value, action_output_dict) = self.evaluate(
                  (action_output, action_output_dict)
              )
  
              distribution_output_value = self.evaluate(
                  _sample_from_distributions(distribution_output)
              )
  
              input_fn_and_spec = None
              if has_input_fn_and_spec:
                  input_fn_and_spec = (
                      _convert_string_vector_to_action_input,
                      tf.TensorSpec((7,), tf.string, name="example"),
                  )
  
              saver = policy_saver.PolicySaver(
                  policy,
                  batch_size=None,
                  use_nest_path_signatures=False,
                  seed=action_seed,
                  input_fn_and_spec=input_fn_and_spec,
                  train_step=global_step,
              )
              path = os.path.join(
                  self.get_temp_dir(), "save_model_action"
              )
              saver.save(path)
  
      with tf.compat.v1.Graph().as_default():
          tf.compat.v1.set_random_seed(self._global_seed)
          with tf.compat.v1.Session().as_default():
              reloaded = tf.compat.v2.saved_model.load(path)
  
              self.assertIn("action", reloaded.signatures)
              reloaded_action = reloaded.signatures["action"]
              if has_input_fn_and_spec:
                  self._compare_input_output_specs(
                      reloaded_action,
                      expected_input_specs=input_fn_and_spec[1],
                      expected_output_spec=policy.policy_step_spec,
                      batch_input=True,
                  )
  
              else:
                  self._compare_input_output_specs(
                      reloaded_action,
                      expected_input_specs=(
                          self._time_step_spec,
                          policy.policy_state_spec,
                      ),
                      expected_output_spec=policy.policy_step_spec,
                      batch_input=True,
                  )
  
              # Reload action_input_values as tensors in the new graph.
              action_input_tensors = tf.nest.map_structure(
                  tf.convert_to_tensor, action_input_values
              )
  
              action_input_spec = (
                  self._time_step_spec,
                  policy.policy_state_spec,
              )
              function_action_input_dict = collections.OrderedDict(
                  (spec.name, value)
                  for (spec, value) in zip(
                      tf.nest.flatten(action_input_spec),
                      tf.nest.flatten(action_input_tensors),
                  )
              )
  
              # NOTE(ebrevdo): The graph-level seeds for the policy and the reloaded
              # model are equal, which in addition to seeding the call to action() and
              # PolicySaver helps ensure equality of the output of action() in both
              # cases.
              self.assertEqual(
                  reloaded_action.graph.seed, self._global_seed
              )
  
              # The seed= argument for the SavedModel action call was given at
              # creation of the PolicySaver.
              if has_input_fn_and_spec:
                  action_string_vector = (
                      _convert_action_input_to_string_vector(
                          action_input_tensors
                      )
                  )
                  action_string_vector_values = self.evaluate(
                      action_string_vector
                  )
                  reloaded_action_output_dict = reloaded_action(
                      action_string_vector
                  )
                  reloaded_action_output = reloaded.action(
                      action_string_vector
                  )
                  reloaded_distribution_output = reloaded.distribution(
                      action_string_vector
                  )
                  self.assertIsInstance(
                      reloaded_distribution_output.action,
                      tfp.distributions.Distribution,
                  )
  
              else:
                  # This is the flat-signature function.
                  reloaded_action_output_dict = reloaded_action(
                      **function_action_input_dict
                  )
                  # This is the non-flat function.
                  reloaded_action_output = reloaded.action(
                      *action_input_tensors
                  )
                  reloaded_distribution_output = reloaded.distribution(
                      *action_input_tensors
                  )
                  self.assertIsInstance(
                      reloaded_distribution_output.action,
                      tfp.distributions.Distribution,
                  )
  
                  if not has_state:
                      # Try both cases: one with an empty policy_state and one with no
                      # policy_state.  Compare them.
  
                      # NOTE(ebrevdo): The first call to .action() must be stored in
                      # reloaded_action_output because this is the version being compared
                      # later against the true action_output and the values will change
                      # after the first call due to randomness.
                      reloaded_action_output_no_input_state = (
                          reloaded.action(action_input_tensors[0])
                      )
                      reloaded_distribution_output_no_input_state = (
                          reloaded.distribution(action_input_tensors[0])
                      )
                      # Even with a seed, multiple calls to action will get different
                      # values, so here we just check the signature matches.
                      self.assertIsInstance(
                          reloaded_distribution_output_no_input_state.action,
                          tfp.distributions.Distribution,
                      )
                      tf.nest.map_structure(
                          self.match_dtype_shape,
                          reloaded_action_output_no_input_state,
                          reloaded_action_output,
                      )
  
                      tf.nest.map_structure(
                          self.match_dtype_shape,
                          _sample_from_distributions(
                              reloaded_distribution_output_no_input_state
                          ),
                          _sample_from_distributions(
                              reloaded_distribution_output
                          ),
                      )
  
              self.evaluate(tf.compat.v1.global_variables_initializer())
              (
                  reloaded_action_output_dict,
                  reloaded_action_output_value,
              ) = self.evaluate(
                  (reloaded_action_output_dict, reloaded_action_output)
              )
  
              reloaded_distribution_output_value = self.evaluate(
                  _sample_from_distributions(
                      reloaded_distribution_output
                  )
              )
  
              self.assertAllEqual(
                  action_output_dict.keys(),
                  reloaded_action_output_dict.keys(),
              )
  
              for k in action_output_dict:
                  if seeded:
                      self.assertAllClose(
                          action_output_dict[k],
                          reloaded_action_output_dict[k],
                          msg="\nMismatched dict key: %s." % k,
                      )
                  else:
                      self.match_dtype_shape(
                          action_output_dict[k],
                          reloaded_action_output_dict[k],
                          msg="\nMismatch dict key: %s." % k,
                      )
  
              # With non-signature functions, we can check that passing a seed does
              # the right thing the second time.
              if seeded:
                  tf.nest.map_structure(
                      self.assertAllClose,
                      action_output_value,
                      reloaded_action_output_value,
                  )
              else:
                  tf.nest.map_structure(
                      self.match_dtype_shape,
                      action_output_value,
                      reloaded_action_output_value,
                  )
  
              tf.nest.map_structure(
                  self.assertAllClose,
                  distribution_output_value,
                  reloaded_distribution_output_value,
              )
  
      ## TFLite tests.
  
      # The converter must run outside of a TF1 graph context, even in
      # eager mode, to ensure the TF2 path is being executed.  Only
      # works in TF2.
      if tf.compat.v1.executing_eagerly_outside_functions():
          tflite_converter = tf.lite.TFLiteConverter.from_saved_model(
              path, signature_keys=["action"]
          )
          # Enable using TF ops in a TFLite model as the TF op `StringToNumber` is
          # not supported natively in TFLite.
          tflite_converter.target_spec.supported_ops = [
              tf.lite.OpsSet.TFLITE_BUILTINS,
              tf.lite.OpsSet.SELECT_TF_OPS,
          ]
          tflite_serialized_model = tflite_converter.convert()
  
          tflite_interpreter = tf.lite.Interpreter(
              model_content=tflite_serialized_model
          )
  
          tflite_runner = tflite_interpreter.get_signature_runner(
              "action"
          )
          tflite_signature = tflite_interpreter.get_signature_list()[
              "action"
          ]
  
          if has_input_fn_and_spec:
              tflite_action_input_dict = {
                  "example": action_string_vector_values,
              }
          else:
              tflite_action_input_dict = collections.OrderedDict(
                  (spec.name, value)
                  for (spec, value) in zip(
                      tf.nest.flatten(action_input_spec),
                      tf.nest.flatten(action_input_values),
                  )
              )
  
          self.assertEqual(
              set(tflite_signature["inputs"]),
              set(tflite_action_input_dict),
          )
          self.assertEqual(
              set(tflite_signature["outputs"]), set(action_output_dict)
          )
  
          if not tf.test.is_gpu_available():
              # TODO(b/191406764): This segfaults due to GPU+Flex ops in
              # Python interpreter.  Allow to run with GPU once bug is fixed.
              tflite_output = tflite_runner(**tflite_action_input_dict)
  
              self.assertAllClose(tflite_output, action_output_dict)
  
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 4 for calling the PolicySaverTest obj's underlying member method code with example usage: obj.testSaveAction(seeded, has_state, distribution_net, has_input_fn_and_spec) without return types