code: def restart(
      observation: types.NestedTensorOrArray,
      batch_size: Optional[types.Int] = None,
      reward_spec: Optional[types.NestedSpec] = None,
  ) -> TimeStep:
      """Returns a `TimeStep` with `step_type` set equal to `StepType.FIRST`.
  
      Args:
        observation: A NumPy array, tensor, or a nested dict, list or tuple of
          arrays or tensors.
        batch_size: (Optional) A python or tensorflow integer scalar. If not
          provided, the environment will not be considered as a batched environment.
        reward_spec: If provided, the reward in the returned `TimeStep` will be
          compatible with the provided `reward_spec`.
  
      Returns:
        A `TimeStep`.
      """
      first_observation = tf.nest.flatten(observation)[0]
      if not tf.is_tensor(first_observation):
          if batch_size is not None:
              if reward_spec is None:
                  reward = np.zeros(batch_size, dtype=np.float32)
              else:
                  reward = tf.nest.map_structure(
                      lambda r: np.zeros(
                          [batch_size] + list(r.shape), _get_np_dtype(r)
                      ),
                      reward_spec,
                  )
              discount = np.ones(batch_size, dtype=np.float32)
              step_type = np.tile(StepType.FIRST, batch_size)
              return TimeStep(step_type, reward, discount, observation)
          else:
              if reward_spec is None:
                  return TimeStep(
                      StepType.FIRST,
                      _as_array(0.0),
                      _as_array(1.0),
                      observation,
                  )
              else:
                  reward = tf.nest.map_structure(
                      lambda r: np.zeros(
                          r.shape, dtype=_get_np_dtype(r)
                      ),
                      reward_spec,
                  )
                  return TimeStep(
                      StepType.FIRST,
                      reward,
                      _as_array(1.0),
                      observation,
                  )
  
      # TODO(b/130244501): Check leading dimension of first_observation
      # against batch_size if all are known statically.
      shape = _as_multi_dim(batch_size)
      step_type = tf.fill(shape, StepType.FIRST, name="step_type")
      if reward_spec is None:
          reward = tf.fill(shape, _as_array(0.0), name="reward")
      else:
          reward = tf.nest.map_structure(
              # pylint: disable=g-long-lambda
              lambda r: tf.fill(
                  tf.concat([shape, r.shape], axis=-1),
                  _as_array(0.0, _get_np_dtype(r)),
                  name="reward",
              ),
              reward_spec,
          )
      discount = tf.fill(shape, _as_array(1.0), name="discount")
      return TimeStep(step_type, reward, discount, observation)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import pprint
  from typing import NamedTuple, Optional
  import numpy
  import tensorflow
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
label_desc: Use image node_id 3 for calling a global function with example usage: restart(observation, batch_size, reward_spec) and returns: TimeStep, TimeStep, TimeStep, TimeStep