code: def _get_initial_and_final_steps(observation_shape, reward):
      batch_size = reward.shape[0]
      time_step_spec = time_step.time_step_spec(
          tensor_spec.TensorSpec(observation_shape, tf.float32)
      )
      initial_step = tensor_spec.sample_spec_nest(
          time_step_spec, outer_dims=(batch_size,)
      )
      final_step = initial_step._replace(
          reward=tf.convert_to_tensor(reward)
      )
      return initial_step, final_step
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import exp3_agent
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 1 for calling a global function with example usage: _get_initial_and_final_steps(observation_shape, reward) and returns: initial_step,  final_step