code: def testSaveRestore(self, batch_size):
      policy_save_path = os.path.join(
          flags.FLAGS.test_tmpdir, "policy", str(batch_size)
      )
  
      # Construct a policy to be saved under a tf.Graph instance.
      policy_saved_graph = tf.Graph()
      with policy_saved_graph.as_default():
          tf_policy = q_policy.QPolicy(
              self._time_step_spec,
              self._action_spec,
              DummyNet(use_constant_initializer=False),
          )
  
          # Parameterized tests reuse temp directories, make no save exists.
          try:
              tf.io.gfile.listdir(policy_save_path)
              tf.io.gfile.rmtree(policy_save_path)
          except tf.errors.NotFoundError:
              pass
          policy_saved = py_tf_policy.PyTFPolicy(tf_policy)
          policy_saved.session = tf.compat.v1.Session(
              graph=policy_saved_graph
          )
          policy_saved.initialize(batch_size)
          policy_saved.save(
              policy_dir=policy_save_path, graph=policy_saved_graph
          )
          # Verify that index files were written. There will also be some number of
          # data files, but this depends on the number of devices.
          self.assertContainsSubset(
              set(["checkpoint", "ckpt-0.index"]),
              set(tf.io.gfile.listdir(policy_save_path)),
          )
  
      # Construct a policy to be restored under another tf.Graph instance.
      policy_restore_graph = tf.Graph()
      with policy_restore_graph.as_default():
          tf_policy = q_policy.QPolicy(
              self._time_step_spec,
              self._action_spec,
              DummyNet(use_constant_initializer=False),
          )
          policy_restored = py_tf_policy.PyTFPolicy(tf_policy)
          policy_restored.session = tf.compat.v1.Session(
              graph=policy_restore_graph
          )
          policy_restored.initialize(batch_size)
          random_init_vals = policy_restored.session.run(
              tf_policy.variables()
          )
          policy_restored.restore(
              policy_dir=policy_save_path, graph=policy_restore_graph
          )
          restored_vals = policy_restored.session.run(
              tf_policy.variables()
          )
          for random_init_var, restored_var in zip(
              random_init_vals, restored_vals
          ):
              self.assertFalse(
                  np.array_equal(random_init_var, restored_var)
              )
  
      # Check that variables in the two policies have identical values.
      with policy_restore_graph.as_default():
          restored_values = policy_restored.session.run(
              tf.compat.v1.global_variables()
          )
      with policy_saved_graph.as_default():
          initial_values = policy_saved.session.run(
              tf.compat.v1.global_variables()
          )
  
      # Networks have two fully connected layers.
      self.assertLen(initial_values, 4)
      self.assertLen(restored_values, 4)
  
      for initial_var, restored_var in zip(
          initial_values, restored_values
      ):
          np.testing.assert_array_equal(initial_var, restored_var)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.networks import network
  from tf_agents.policies import actor_policy
  from tf_agents.policies import ou_noise_policy
  from tf_agents.policies import py_tf_policy
  from tf_agents.policies import q_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 7 for calling the PyTFPolicyTest obj's underlying member method code with example usage: obj.testSaveRestore(batch_size) without return types