code: def testInferenceWithCheckpoint(self):
      if not common.has_eager_been_enabled():
          self.skipTest("Only supported in TF2.x.")
  
      # Create and saved_model for a q_policy.
      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,
      )
      sample_input = self.evaluate(
          tensor_spec.sample_spec_nest(
              self._time_step_spec, outer_dims=(3,)
          )
      )
  
      saver = policy_saver.PolicySaver(policy, batch_size=None)
      path = os.path.join(self.get_temp_dir(), "save_model")
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      original_eval = self.evaluate(policy.action(sample_input))
      saver.save(path)
      # Asign -1 to all variables in the policy. Making checkpoint different than
      # the initial saved_model.
      self.evaluate(
          tf.nest.map_structure(
              lambda v: v.assign(v * 0 + -1), policy.variables()
          )
      )
      checkpoint_path = os.path.join(self.get_temp_dir(), "checkpoint")
      saver.save_checkpoint(checkpoint_path)
  
      # Get an instance of the saved_model.
      reloaded_policy = tf.compat.v2.saved_model.load(path)
      self.evaluate(
          tf.compat.v1.initializers.variables(
              reloaded_policy.model_variables
          )
      )
  
      # Verify loaded saved_model variables are different than the current policy.
      model_variables = self.evaluate(policy.variables())
      reloaded_model_variables = self.evaluate(
          reloaded_policy.model_variables
      )
  
      any_not_equal = lambda a, b: np.not_equal(a, b).any()
      self.assertTrue(
          any(
              [
                  any_not_equal(a, b)
                  for a, b in zip(
                      tf.nest.flatten(model_variables),
                      tf.nest.flatten(reloaded_model_variables),
                  )
              ]
          )
      )
  
      # Update from checkpoint.
      checkpoint = tf.train.Checkpoint(policy=reloaded_policy)
      checkpoint_file_prefix = os.path.join(
          checkpoint_path, "variables", "variables"
      )
      checkpoint.read(
          checkpoint_file_prefix
      ).assert_existing_objects_matched()
  
      self.evaluate(
          tf.compat.v1.initializers.variables(
              reloaded_policy.model_variables
          )
      )
  
      # Verify variables are now equal.
      model_variables = self.evaluate(policy.variables())
      reloaded_model_variables = self.evaluate(
          reloaded_policy.model_variables
      )
  
      all_equal = lambda a, b: np.equal(a, b).all()
      self.assertTrue(
          all(
              [
                  all_equal(a, b)
                  for a, b in zip(
                      tf.nest.flatten(model_variables),
                      tf.nest.flatten(reloaded_model_variables),
                  )
              ]
          )
      )
  
      # Verify variable update affects inference.
      assert_np_not_equal = lambda a, b: self.assertFalse(
          np.equal(a, b).all()
      )
      reloaded_eval = self.evaluate(
          reloaded_policy.action(sample_input)
      )
      tf.nest.map_structure(
          assert_np_not_equal, original_eval, reloaded_eval
      )
      current_eval = self.evaluate(policy.action(sample_input))
      tf.nest.map_structure(
          assert_np_not_equal, current_eval, reloaded_eval
      )
  
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 18 for calling the PolicySaverTest obj's underlying member method code with example usage: obj.testInferenceWithCheckpoint() without return types