code: def trajectory_for_bandit(
      initial_step: types.TimeStep,
      action_step: types.PolicyStep,
      final_step: types.TimeStep,
  ) -> types.NestedTensor:
      """Builds a trajectory from a single-step bandit episode.
  
      Since all episodes consist of a single step, the returned `Trajectory` has no
      time dimension. All input and output `Tensor`s/arrays are expected to have
      shape `[batch_size, ...]`.
  
      Args:
        initial_step: A `TimeStep` returned from `environment.step(...)`.
        action_step: A `PolicyStep` returned by `policy.action(...)`.
        final_step: A `TimeStep` returned from `environment.step(...)`.
  
      Returns:
        A `Trajectory` containing zeros for discount value and `StepType.LAST` for
        both `step_type` and `next_step_type`.
      """
      return trajectory.Trajectory(
          observation=initial_step.observation,
          action=action_step.action,
          policy_info=action_step.info,
          reward=final_step.reward,
          discount=final_step.discount,
          step_type=initial_step.step_type,
          next_step_type=final_step.step_type,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
label_desc: Use image node_id 1 for calling a global function with example usage: trajectory_for_bandit(initial_step, action_step, final_step) and returns: trajectory