docstr: Helper function for testing correct variable updating and restoring.
docstr_tok: ["Helper","function","for","testing","correct","variable","updating","and","restoring","."]
code: def verifyTrainAndRestore(
      self, observation_spec, action_spec, actor_net, loss_fn=None
  ):
      """Helper function for testing correct variable updating and restoring."""
      batch_size = 2
      observations = tensor_spec.sample_spec_nest(
          observation_spec, outer_dims=(batch_size,)
      )
      actions = tensor_spec.sample_spec_nest(
          action_spec, outer_dims=(batch_size,)
      )
      rewards = tf.constant([10, 20], dtype=tf.float32)
      discounts = tf.constant([0.9, 0.9], dtype=tf.float32)
      experience = trajectory.first(
          observation=observations,
          action=actions,
          policy_info=(),
          reward=rewards,
          discount=discounts,
      )
      time_step_spec = ts.time_step_spec(observation_spec)
      strategy = tf.distribute.get_strategy()
      with strategy.scope():
          # Use BehaviorCloningAgent instead of AWRAgent to test the network.
          agent = behavioral_cloning_agent.BehavioralCloningAgent(
              time_step_spec,
              action_spec,
              cloning_network=actor_net,
              optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
              loss_fn=loss_fn,
          )
      loss_before_train = agent.loss(experience).loss
      # Check loss is stable.
      self.assertEqual(loss_before_train, agent.loss(experience).loss)
      # Train 2 steps, verify that loss is decreased for the same input.
      for _ in range(2):
          agent.train(experience)
      loss_after_train = agent.loss(experience).loss
      self.assertLessEqual(loss_after_train, loss_before_train)
      # Assert loss evaluation is still stable, e.g. deterministic.
      self.assertLessEqual(
          loss_after_train, agent.loss(experience).loss
      )
      # Save checkpoint
      ckpt_dir = self.create_tempdir()
      checkpointer = common.Checkpointer(ckpt_dir=ckpt_dir, agent=agent)
      global_step = tf.constant(1)
      checkpointer.save(global_step)
      # Assign all vars to 0.
      for var in tf.nest.flatten(agent.variables):
          var.assign(tf.zeros_like(var))
      loss_after_zero = agent.loss(experience).loss
      self.assertEqual(loss_after_zero, agent.loss(experience).loss)
      self.assertNotEqual(loss_after_zero, loss_after_train)
      # Restore
      checkpointer._checkpoint.restore(
          checkpointer._manager.latest_checkpoint
      )
      loss_after_restore = agent.loss(experience).loss
      self.assertNotEqual(loss_after_restore, loss_after_zero)
      self.assertEqual(loss_after_restore, loss_after_train)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents import test_util
  from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
  from tf_agents.keras_layers import inner_reshape
  from tf_agents.networks import expand_dims_layer
  from tf_agents.networks import q_network
  from tf_agents.networks import q_rnn_network
  from tf_agents.networks import sequential
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 5 for calling the BehavioralCloningAgentTest obj's underlying member method code with example usage: obj.verifyTrainAndRestore(observation_spec, action_spec, actor_net, loss_fn) without return types