code: def testMetricIsComputedCorrectly(
      self, num_time_steps, expected_reward
  ):
      batched_avg_return_metric = batched_py_metric.BatchedPyMetric(
          py_metrics.AverageReturnMetric
      )
      tf_avg_return_metric = tf_py_metric.TFPyMetric(
          batched_avg_return_metric
      )
      deps = []
      for i in range(num_time_steps):
          with tf.control_dependencies(deps):
              traj = tf_avg_return_metric(self._ts[i])
              deps = tf.nest.flatten(traj)
      with tf.control_dependencies(deps):
          result = tf_avg_return_metric.result()
      result_ = self.evaluate(result)
      self.assertEqual(result_, expected_reward)
  
file_path: agents/tf_agents/metrics/tf_py_metric_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  from tf_agents.metrics import batched_py_metric
  from tf_agents.metrics import py_metrics
  from tf_agents.metrics import tf_py_metric
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 2 for calling the BatchedPyMetricTest obj's underlying member method code with example usage: obj.testMetricIsComputedCorrectly(num_time_steps, expected_reward) without return types