code: def _create_trajectories(
      n_time_steps, batch_size, policy_info_scale_name
  ):
      # Observation looks like:
      # [[ 0.,  1., ... n_time_steps.],
      #  [10., 11., ... n_time_steps.],
      #  [20., 21., ... n_time_steps.],
      #  [ ...                       ],
      #  [10*batch_size., ... 10*batch_size+n_time_steps.]]
      observation_array = np.asarray(
          [np.arange(n_time_steps) + 10 * i for i in range(batch_size)]
      )
      # Adding an inner most dimension to fit the observation spec defined above.
      observation_array = np.expand_dims(observation_array, axis=2)
      observations = tf.convert_to_tensor(
          observation_array, dtype=tf.float32
      )
  
      default_tensor = tf.constant(
          [[1] * n_time_steps] * batch_size, dtype=tf.float32
      )
      mid_time_step_val = ts.StepType.MID.tolist()
      time_steps = ts.TimeStep(
          step_type=tf.constant(
              [[mid_time_step_val] * n_time_steps] * batch_size,
              dtype=tf.int32,
          ),
          reward=default_tensor,
          discount=default_tensor,
          observation=observations,
      )
      actions = tf.constant(
          [[[1]] * n_time_steps] * batch_size, dtype=tf.float32
      )
      policy_info = {
          "dist_params": {
              "loc": tf.constant(
                  [[[1]] * n_time_steps] * batch_size, dtype=tf.float32
              ),
              policy_info_scale_name: tf.constant(
                  [[[1]] * n_time_steps] * batch_size, dtype=tf.float32
              ),
          },
          "value_prediction": default_tensor,
          "return": default_tensor,
          "advantage": default_tensor,
      }
      return trajectory.Trajectory(
          time_steps.step_type,
          observations,
          actions,
          policy_info,
          time_steps.step_type,
          time_steps.reward,
          time_steps.discount,
      )
  
imports: from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.agents.ppo import ppo_actor_network
  from tf_agents.agents.ppo import ppo_agent
  from tf_agents.agents.ppo import ppo_policy
  from tf_agents.distributions import utils
  from tf_agents.networks import actor_distribution_network
  from tf_agents.networks import value_network
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _create_trajectories(n_time_steps, batch_size, policy_info_scale_name) and returns: trajectory