Record the information obtained from taking a step in all envs.

    Records (observation, rewards, done) in a new time-step and actions in the
    current time-step.

    If any trajectory gets done, we move that trajectory to
    completed_trajectories.

    Args:
      observations: ndarray of first dimension self.batch_size, which has the
        observations after we've stepped, i.e. s_{t+1} where t is the current
        state.
      raw_rewards: ndarray of first dimension self.batch_size containing raw
        rewards i.e. r_{t+1}.
      processed_rewards: ndarray of first dimension self.batch_size containing
        processed rewards. i.e. r_{t+1}
      dones: ndarray of first dimension self.batch_size, containing true at an
        index if that env is done, i.e. d_{t+1}
      actions: ndarray of first dimension self.batch_size, containing actions
        applied at the current time-step, which leads to the observations
        rewards and done at the next time-step, i.e. a_t

def step(self, observations, raw_rewards, processed_rewards, dones, actions):
    """Record the information obtained from taking a step in all envs.

    Records (observation, rewards, done) in a new time-step and actions in the
    current time-step.

    If any trajectory gets done, we move that trajectory to
    completed_trajectories.

    Args:
      observations: ndarray of first dimension self.batch_size, which has the
        observations after we've stepped, i.e. s_{t+1} where t is the current
        state.
      raw_rewards: ndarray of first dimension self.batch_size containing raw
        rewards i.e. r_{t+1}.
      processed_rewards: ndarray of first dimension self.batch_size containing
        processed rewards. i.e. r_{t+1}
      dones: ndarray of first dimension self.batch_size, containing true at an
        index if that env is done, i.e. d_{t+1}
      actions: ndarray of first dimension self.batch_size, containing actions
        applied at the current time-step, which leads to the observations
        rewards and done at the next time-step, i.e. a_t
    """
    # Pre-conditions
    assert isinstance(observations, np.ndarray)
    assert isinstance(raw_rewards, np.ndarray)
    assert isinstance(processed_rewards, np.ndarray)
    assert isinstance(dones, np.ndarray)
    assert isinstance(actions, np.ndarray)

    # We assume that we step in all envs, i.e. not like reset where we can reset
    # some envs and not others.
    assert self.batch_size == observations.shape[0]
    assert self.batch_size == raw_rewards.shape[0]
    assert self.batch_size == processed_rewards.shape[0]
    assert self.batch_size == dones.shape[0]
    assert self.batch_size == actions.shape[0]

    for index in range(self.batch_size):
      trajectory = self._trajectories[index]

      # NOTE: If the trajectory isn't active, that means it doesn't have any
      # time-steps in it, but we are in step, so the assumption is that it has
      # a prior observation from which we are stepping away from.

      # TODO(afrozm): Let's re-visit this if it becomes too restrictive.
      assert trajectory.is_active

      # To this trajectory's last time-step, set actions.
      trajectory.change_last_time_step(action=actions[index])

      # Create a new time-step to add observation, done & rewards (no actions).
      trajectory.add_time_step(
          observation=observations[index],
          done=dones[index],
          raw_reward=raw_rewards[index],
          processed_reward=processed_rewards[index])

      # If the trajectory is completed, i.e. dones[index] == True, then we
      # account for it right-away.
      if dones[index]:
        self._complete_trajectory(trajectory, index)

        # NOTE: The new trajectory at `index` is going to be in-active and
        # `reset` should be called on it.
        assert not self._trajectories[index].is_active