<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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