code: def testBatchRestart(self):
      obs_spec = [tensor_spec.TensorSpec([2], tf.float32)]
      time_step_spec = ts.time_step_spec(obs_spec)
      observations = [tf.constant([[1, 2], [3, 4]], dtype=tf.float32)]
      time_steps = ts.restart(observations, 2)
      time_step_batched = nest_utils.is_batched_nested_tensors(
          time_steps, time_step_spec
      )
      self.assertTrue(time_step_batched)
  
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 2 for calling the TFTimeStepTest obj's underlying member method code with example usage: obj.testBatchRestart() without return types