code: def _get_initial_and_final_steps(observations, scores):
      batch_size = tf.nest.flatten(observations)[0].shape[0]
      initial_step = ts.TimeStep(
          tf.constant(
              ts.StepType.FIRST,
              dtype=tf.int32,
              shape=[batch_size],
              name="step_type",
          ),
          tf.constant(
              0.0, dtype=tf.float32, shape=[batch_size], name="reward"
          ),
          tf.constant(
              1.0, dtype=tf.float32, shape=[batch_size], name="discount"
          ),
          observations,
      )
      final_step = ts.TimeStep(
          tf.constant(
              ts.StepType.LAST,
              dtype=tf.int32,
              shape=[batch_size],
              name="step_type",
          ),
          scores,
          tf.constant(
              1.0, dtype=tf.float32, shape=[batch_size], name="discount"
          ),
          observations,
      )
      return initial_step, final_step
  
file_path: agents/tf_agents/bandits/agents/ranking_agent_test.py
imports: from absl.testing import parameterized
  import tensorflow
  from tf_agents.bandits.agents import ranking_agent
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.policies import utils
  from tf_agents.specs import bandit_spec_utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _get_initial_and_final_steps(observations, scores) and returns: initial_step,  final_step