Collect trajectories with the given policy net and behaviour.

  Args:
    env: A gym env interface, for now this is not-batched.
    policy_fun: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
    num_trajectories: int, number of trajectories.
    policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
        how to use the policy_fun to return an action.
    max_timestep: int or None, the index of the maximum time-step at which we
        return the trajectory, None for ending a trajectory only when env
        returns done.
    epsilon: float, the epsilon for `epsilon-greedy` policy.

  Returns:
    trajectory: list of (observation, action, reward) tuples, where each element
    `i` is a tuple of numpy arrays with shapes as follows:
    observation[i] = (B, T_i + 1)
    action[i] = (B, T_i)
    reward[i] = (B, T_i)

def collect_trajectories(env,
                         policy_fun,
                         num_trajectories=1,
                         policy="greedy",
                         max_timestep=None,
                         epsilon=0.1):
  """Collect trajectories with the given policy net and behaviour.

  Args:
    env: A gym env interface, for now this is not-batched.
    policy_fun: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
    num_trajectories: int, number of trajectories.
    policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
        how to use the policy_fun to return an action.
    max_timestep: int or None, the index of the maximum time-step at which we
        return the trajectory, None for ending a trajectory only when env
        returns done.
    epsilon: float, the epsilon for `epsilon-greedy` policy.

  Returns:
    trajectory: list of (observation, action, reward) tuples, where each element
    `i` is a tuple of numpy arrays with shapes as follows:
    observation[i] = (B, T_i + 1)
    action[i] = (B, T_i)
    reward[i] = (B, T_i)
  """
  trajectories = []

  for t in range(num_trajectories):
    t_start = time.time()
    rewards = []
    actions = []
    done = False

    observation = env.reset()

    # This is currently shaped (1, 1) + OBS, but new observations will keep
    # getting added to it, making it eventually (1, T+1) + OBS
    observation_history = observation[np.newaxis, np.newaxis, :]

    # Run either till we're done OR if max_timestep is defined only till that
    # timestep.
    ts = 0
    while ((not done) and
           (not max_timestep or observation_history.shape[1] < max_timestep)):
      ts_start = time.time()
      # Run the policy, to pick an action, shape is (1, t, A) because
      # observation_history is shaped (1, t) + OBS
      predictions = policy_fun(observation_history)

      # We need the predictions for the last time-step, so squeeze the batch
      # dimension and take the last time-step.
      predictions = np.squeeze(predictions, axis=0)[-1]

      # Policy can be run in one of the following ways:
      #  - Greedy
      #  - Epsilon-Greedy
      #  - Categorical-Sampling
      action = None
      if policy == "greedy":
        action = np.argmax(predictions)
      elif policy == "epsilon-greedy":
        # A schedule for epsilon is 1/k where k is the episode number sampled.
        if onp.random.random() < epsilon:
          # Choose an action at random.
          action = onp.random.randint(0, high=len(predictions))
        else:
          # Return the best action.
          action = np.argmax(predictions)
      elif policy == "categorical-sampling":
        # NOTE: The predictions aren't probabilities but log-probabilities
        # instead, since they were computed with LogSoftmax.
        # So just np.exp them to make them probabilities.
        predictions = np.exp(predictions)
        action = onp.argwhere(onp.random.multinomial(1, predictions) == 1)
      else:
        raise ValueError("Unknown policy: %s" % policy)

      # NOTE: Assumption, single batch.
      try:
        action = int(action)
      except TypeError as err:
        # Let's dump some information before we die off.
        logging.error("Cannot convert action into an integer: [%s]", err)
        logging.error("action.shape: [%s]", action.shape)
        logging.error("action: [%s]", action)
        logging.error("predictions.shape: [%s]", predictions.shape)
        logging.error("predictions: [%s]", predictions)
        logging.error("observation_history: [%s]", observation_history)
        raise err

      observation, reward, done, _ = env.step(action)

      # observation is of shape OBS, so add extra dims and concatenate on the
      # time dimension.
      observation_history = np.concatenate(
          [observation_history, observation[np.newaxis, np.newaxis, :]], axis=1)

      rewards.append(reward)
      actions.append(action)

      ts += 1
      logging.vlog(
          2, "  Collected time-step[ %5d] of trajectory[ %5d] in [%0.2f] msec.",
          ts, t, get_time(ts_start))
    logging.vlog(
        2, " Collected trajectory[ %5d] in [%0.2f] msec.", t, get_time(t_start))

    # This means we are done we're been terminated early.
    assert done or (
        max_timestep and max_timestep >= observation_history.shape[1])
    # observation_history is (1, T+1) + OBS, lets squeeze out the batch dim.
    observation_history = np.squeeze(observation_history, axis=0)
    trajectories.append(
        (observation_history, np.stack(actions), np.stack(rewards)))

  return trajectories