code: def traj_segment_generator(
      policy, env, horizon, reward_giver=None, gail=False, callback=None
  ):
      """
      Compute target value using TD(lambda) estimator, and advantage with GAE(lambda)
      :param policy: (MLPPolicy) the policy
      :param env: (Gym Environment) the environment
      :param horizon: (int) the number of timesteps to run per batch
      :param reward_giver: (TransitionClassifier) the reward predicter from obsevation and action
      :param gail: (bool) Whether we are using this generator for standard trpo or with gail
      :param callback: (BaseCallback)
      :return: (dict) generator that returns a dict with the following keys:
          - observations: (np.ndarray) observations
          - rewards: (numpy float) rewards (if gail is used it is the predicted reward)
          - true_rewards: (numpy float) if gail is used it is the original reward
          - vpred: (numpy float) action logits
          - dones: (numpy bool) dones (is end of episode, used for logging)
          - episode_starts: (numpy bool)
              True if first timestep of an episode, used for GAE
          - actions: (np.ndarray) actions
          - nextvpred: (numpy float) next action logits
          - ep_rets: (float) cumulated current episode reward
          - ep_lens: (int) the length of the current episode
          - ep_true_rets: (float) the real environment reward
          - continue_training: (bool) Whether to continue training
              or stop early (triggered by the callback)
      """
      # Check when using GAIL
      assert not (
          gail and reward_giver is None
      ), "You must pass a reward giver when using GAIL"
  
      # Initialize state variables
      step = 0
      action = (
          env.action_space.sample()
      )  # not used, just so we have the datatype
      observation = env.reset()
  
      cur_ep_ret = 0  # return in current episode
      current_it_len = 0  # len of current iteration
      current_ep_len = 0  # len of current episode
      cur_ep_true_ret = 0
      ep_true_rets = []
      ep_rets = []  # returns of completed episodes in this segment
      ep_lens = []  # Episode lengths
  
      # Initialize history arrays
      observations = np.array([observation for _ in range(horizon)])
      true_rewards = np.zeros(horizon, "float32")
      rewards = np.zeros(horizon, "float32")
      vpreds = np.zeros(horizon, "float32")
      episode_starts = np.zeros(horizon, "bool")
      dones = np.zeros(horizon, "bool")
      actions = np.array([action for _ in range(horizon)])
      states = policy.initial_state
      episode_start = (
          True  # marks if we're on first timestep of an episode
      )
      done = False
  
      callback.on_rollout_start()
  
      while True:
          action, vpred, states, _ = policy.step(
              observation.reshape(-1, *observation.shape), states, done
          )
          # Slight weirdness here because we need value function at time T
          # before returning segment [0, T-1] so we get the correct
          # terminal value
          if step > 0 and step % horizon == 0:
              callback.update_locals(locals())
              callback.on_rollout_end()
              yield {
                  "observations": observations,
                  "rewards": rewards,
                  "dones": dones,
                  "episode_starts": episode_starts,
                  "true_rewards": true_rewards,
                  "vpred": vpreds,
                  "actions": actions,
                  "nextvpred": vpred[0] * (1 - episode_start),
                  "ep_rets": ep_rets,
                  "ep_lens": ep_lens,
                  "ep_true_rets": ep_true_rets,
                  "total_timestep": current_it_len,
                  "continue_training": True,
              }
              _, vpred, _, _ = policy.step(
                  observation.reshape(-1, *observation.shape)
              )
              # Be careful!!! if you change the downstream algorithm to aggregate
              # several of these batches, then be sure to do a deepcopy
              ep_rets = []
              ep_true_rets = []
              ep_lens = []
              # Reset current iteration length
              current_it_len = 0
              callback.on_rollout_start()
          i = step % horizon
          observations[i] = observation
          vpreds[i] = vpred[0]
          actions[i] = action[0]
          episode_starts[i] = episode_start
  
          clipped_action = action
          # Clip the actions to avoid out of bound error
          if isinstance(env.action_space, gym.spaces.Box):
              clipped_action = np.clip(
                  action, env.action_space.low, env.action_space.high
              )
  
          if gail:
              reward = reward_giver.get_reward(
                  observation, clipped_action[0]
              )
              observation, true_reward, done, info = env.step(
                  clipped_action[0]
              )
          else:
              observation, reward, done, info = env.step(
                  clipped_action[0]
              )
              true_reward = reward
  
          if callback is not None:
              callback.update_locals(locals())
              if callback.on_step() is False:
                  # We have to return everything so pytype does not complain
                  yield {
                      "observations": observations,
                      "rewards": rewards,
                      "dones": dones,
                      "episode_starts": episode_starts,
                      "true_rewards": true_rewards,
                      "vpred": vpreds,
                      "actions": actions,
                      "nextvpred": vpred[0] * (1 - episode_start),
                      "ep_rets": ep_rets,
                      "ep_lens": ep_lens,
                      "ep_true_rets": ep_true_rets,
                      "total_timestep": current_it_len,
                      "continue_training": False,
                  }
                  return
  
          rewards[i] = reward
          true_rewards[i] = true_reward
          dones[i] = done
          episode_start = done
  
          cur_ep_ret += reward
          cur_ep_true_ret += true_reward
          current_it_len += 1
          current_ep_len += 1
          if done:
              # Retrieve unnormalized reward if using Monitor wrapper
              maybe_ep_info = info.get("episode")
              if maybe_ep_info is not None:
                  if not gail:
                      cur_ep_ret = maybe_ep_info["r"]
                  cur_ep_true_ret = maybe_ep_info["r"]
  
              ep_rets.append(cur_ep_ret)
              ep_true_rets.append(cur_ep_true_ret)
              ep_lens.append(current_ep_len)
              cur_ep_ret = 0
              cur_ep_true_ret = 0
              current_ep_len = 0
              if not isinstance(env, VecEnv):
                  observation = env.reset()
          step += 1
  
file_path: stable-baselines/stable_baselines/common/runners.py
imports: from abc import ABC, abstractmethod
  import typing
  from typing import Union, Optional, Any
  import gym
  import numpy
  from stable_baselines.common.callbacks import BaseCallback
  from stable_baselines.common.vec_env import VecEnv
label_desc: Use image node_id 1 for calling a global function with example usage: traj_segment_generator(policy, env, horizon, reward_giver, gail, callback) without return types