class_docstr: Basically a batch of active trajectories and a list of completed ones.
class_docstr_tok: ["Basically","a","batch","of","active","trajectories","and","a","list","of","completed","ones","."]
docstr: 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
    infos: (optional) a dictionary of keys and values, where all the values
      have the first dimension as self.batch_size.
code: def step(
      self,
      observations,
      raw_rewards,
      processed_rewards,
      dones,
      actions,
      infos=None,
  ):
      """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
        infos: (optional) a dictionary of keys and values, where all the values
          have the first dimension as self.batch_size.
      """
      # 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)
      if infos:
          assert isinstance(infos, dict)
  
      # 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]
      if infos:
          for _, v in infos.items():
              assert self.batch_size == len(v)
  
      def extract_info_at_index(infos, index):
          if not infos:
              return None
          return {k: v[index] for k, v in infos.items()}
  
      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],
              info=extract_info_at_index(infos, 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import pickle
  import re
  import sys
  import time
  from absl import logging
  import cloudpickle
  import numpy
  from tensor2tensor.envs import time_step
  import tensorflow.compat.v1
label_desc: Use image node_id 11 for calling the BatchTrajectory obj's underlying member method code with example usage: obj.step(observations, raw_rewards, processed_rewards, dones, actions, infos) and returns: unknown, None