code: def testTerminationMultiRewards(self):
      observation = np.array([[-1], [-1]])
      reward = [
          np.array([[2.0], [2.0]]),
          np.array([[3.0, 3.0], [4.0, 4.0]]),
      ]
      time_step = ts.termination(observation, reward)
  
      self.assertItemsEqual([ts.StepType.LAST] * 2, time_step.step_type)
      self.assertItemsEqual(observation, time_step.observation)
      self.assertAllEqual(reward[0], time_step.reward[0])
      self.assertAllEqual(reward[1], time_step.reward[1])
      self.assertItemsEqual([0.0, 0.0], time_step.discount)
  
      reward = np.array([[2.0, 2.0, 2.0], [3.0, 3.0, 3.0]])
      reward_spec = [array_spec.ArraySpec((3,), np.float32, "multi_r")]
      outer_dims = nest_utils.get_outer_array_shape(reward, reward_spec)
      time_step_batch = ts.termination(observation, reward, outer_dims)
  
      # Check that passing outer_dims works
      self.assertItemsEqual(
          [ts.StepType.LAST] * 2, time_step_batch.step_type
      )
      self.assertItemsEqual(observation, time_step_batch.observation)
      self.assertAllEqual(reward[0], time_step_batch.reward[0])
      self.assertAllEqual(reward[1], time_step_batch.reward[1])
      self.assertItemsEqual([0.0, 0.0], time_step_batch.discount)
  
      # Check that it gets a different result with no outer_dims
      time_step_no_batch = ts.termination(
          observation, reward, outer_dims=[]
      )
      self.assertEqual(ts.StepType.LAST, time_step_no_batch.step_type)
      self.assertItemsEqual(observation, time_step_no_batch.observation)
      self.assertAllEqual(reward[0], time_step_no_batch.reward[0])
      self.assertAllEqual(reward[1], time_step_no_batch.reward[1])
      self.assertEqual(0.0, time_step_no_batch.discount)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  import numpy
  import tensorflow
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 19 for calling the TimeStepTest obj's underlying member method code with example usage: obj.testTerminationMultiRewards() without return types