code: def test_nested_kl_divergence(self):
      zero = tf.constant([0.0] * 3, dtype=tf.float32)
      one = tf.constant([1.0] * 3, dtype=tf.float32)
      dist_neg_one = tfp.distributions.Normal(loc=-one, scale=one)
      dist_zero = tfp.distributions.Normal(loc=zero, scale=one)
      dist_one = tfp.distributions.Normal(loc=one, scale=one)
  
      nested_dist1 = [dist_zero, [dist_neg_one, dist_one]]
      nested_dist2 = [dist_one, [dist_one, dist_zero]]
      kl_divergence = ppo_utils.nested_kl_divergence(
          nested_dist1, nested_dist2
      )
      expected_kl_divergence = 3 * 3.0  # 3 * (0.5 + (2.0 + 0.5))
  
      kl_divergence_ = self.evaluate(kl_divergence)
      self.assertAllClose(expected_kl_divergence, kl_divergence_)
  
      # test for distributions with different shapes
      one_reshaped = tf.constant([[1.0]] * 3, dtype=tf.float32)
      dist_neg_one_reshaped = tfp.distributions.Normal(
          loc=-one_reshaped, scale=one_reshaped
      )
      dist_one_reshaped = tfp.distributions.Normal(
          loc=one_reshaped, scale=one_reshaped
      )
  
      nested_dist1 = [dist_zero, [dist_neg_one_reshaped, dist_one]]
      nested_dist2 = [dist_one, [dist_one_reshaped, dist_zero]]
      kl_divergence = ppo_utils.nested_kl_divergence(
          nested_dist1, nested_dist2
      )
      expected_kl_divergence = 3 * 3.0  # 3 * (0.5 + (2.0 + 0.5))
  
      kl_divergence_ = self.evaluate(kl_divergence)
      self.assertAllClose(expected_kl_divergence, kl_divergence_)
  
file_path: agents/tf_agents/agents/ppo/ppo_utils_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents.ppo import ppo_utils
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 2 for calling the PPOUtilsTest obj's underlying member method code with example usage: obj.test_nested_kl_divergence() without return types