class_docstr: Collects samples for one agent in one trajectory (episode).
  
  The agent may be part of a multi-agent environment. Samples are stored in
  lists including some possible automatic "shift" buffer at the beginning to
  be able to save memory when storing things like NEXT_OBS, PREV_REWARDS,
  etc.., which are specified using the trajectory view API.
class_docstr_tok: ["Collects","samples","for","one","agent","in","one","trajectory","(","episode",")",".","The","agent","may","be","part","of","a","multi-agent","environment",".","Samples","are","stored","in","lists","including","some","possible","automatic","``","shift","''","buffer","at","the","beginning","to","be","able","to","save","memory","when","storing","things","like","NEXT_OBS",",","PREV_REWARDS",",","etc","..",",","which","are","specified","using","the","trajectory","view","API","."]
docstr: Builds a SampleBatch from the thus-far collected agent data.
  
  If the episode/trajectory has no TERMINATED|TRUNCATED=True at the end, will
  copy the necessary n timesteps at the end of the trajectory back to the
  beginning of the buffers and wait for new samples coming in.
  SampleBatches created by this method will be ready for postprocessing
  by a Policy.
  
  Args:
      view_requirements: The viewrequirements dict needed to build the
      SampleBatch from the raw buffers (which may have data shifts as well as
      mappings from view-col to data-col in them).
  
  Returns:
      SampleBatch: The built SampleBatch for this agent, ready to go into
      postprocessing.
docstr_tok: ["Builds","a","SampleBatch","from","the","thus-far","collected","agent","data",".","If","the","episode\/trajectory","has","no","TERMINATED|TRUNCATED=True","at","the","end",",","will","copy","the","necessary","n","timesteps","at","the","end","of","the","trajectory","back","to","the","beginning","of","the","buffers","and","wait","for","new","samples","coming","in",".","SampleBatches","created","by","this","method","will","be","ready","for","postprocessing","by","a","Policy",".","Args",":","view_requirements",":","The","viewrequirements","dict","needed","to","build","the","SampleBatch","from","the","raw","buffers","(","which","may","have","data","shifts","as","well","as","mappings","from","view-col","to","data-col","in","them",")",".","Returns",":","SampleBatch",":","The","built","SampleBatch","for","this","agent",",","ready","to","go","into","postprocessing","."]
code: def build_for_training(
      self, view_requirements: ViewRequirementsDict
  ) -> SampleBatch:
      """Builds a SampleBatch from the thus-far collected agent data.
  
      If the episode/trajectory has no TERMINATED|TRUNCATED=True at the end, will
      copy the necessary n timesteps at the end of the trajectory back to the
      beginning of the buffers and wait for new samples coming in.
      SampleBatches created by this method will be ready for postprocessing
      by a Policy.
  
      Args:
          view_requirements: The viewrequirements dict needed to build the
          SampleBatch from the raw buffers (which may have data shifts as well as
          mappings from view-col to data-col in them).
  
      Returns:
          SampleBatch: The built SampleBatch for this agent, ready to go into
          postprocessing.
      """
      batch_data = {}
      np_data = {}
      for view_col, view_req in view_requirements.items():
          # Create the batch of data from the different buffers.
          data_col = view_req.data_col or view_col
  
          if data_col not in self.buffers:
              is_state = self._fill_buffer_with_initial_values(
                  data_col, view_req, build_for_inference=False
              )
  
              # We need to skip this view_col if it does not exist in the buffers and
              # is not an RNN state because it could be the special keys that gets
              # added by policy's postprocessing function for training.
              if not is_state:
                  continue
  
          # OBS and INFOS are already shifted by -1 (the initial obs/info starts one
          # ts before all other data columns).
          obs_shift = (
              -1
              if data_col in [SampleBatch.OBS, SampleBatch.INFOS]
              else 0
          )
  
          # Keep an np-array cache so we don't have to regenerate the
          # np-array for different view_cols using to the same data_col.
          self._cache_in_np(np_data, data_col)
  
          # Go through each time-step in the buffer and construct the view
          # accordingly.
          data = []
          for d in np_data[data_col]:
              shifted_data = []
  
              # batch_repeat_value determines how many time steps should we skip
              # before we repeat indexing the data.
              # Example: batch_repeat_value=10, shift_arr = [-3, -2, -1],
              # shift_before = 3
              # buffer = [-3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
              # resulting_data = [[-3, -2, -1], [7, 8, 9]]
              # explanation: For t=0, we output [-3, -2, -1]. We then skip 10 time
              # steps ahead and get to t=10. For t=10, we output [7, 8, 9]. We skip
              # 10 more time steps and get to t=20. but since t=20 is out of bound we
              # stop.
  
              # count computes the number of time steps that we need to consider.
              # if batch_repeat_value = 1, this number should be the length of
              # episode so far, which is len(buffer) - shift_before (-1 if this
              # value was gained during inference. This is because we keep a dummy
              # value at the last position of the buffer that makes it one longer).
              count = int(
                  math.ceil(
                      (
                          len(d)
                          - int(
                              data_col
                              in self.data_cols_with_dummy_values
                          )
                          - self.shift_before
                      )
                      / view_req.batch_repeat_value
                  )
              )
              for i in range(count):
                  # the indices for time step t
                  inds = (
                      self.shift_before
                      + obs_shift
                      + view_req.shift_arr
                      + (i * view_req.batch_repeat_value)
                  )
  
                  # handle the case where the inds are out of bounds from the end.
                  # if during the indexing any of the indices are out of bounds, we
                  # need to use padding on the end to fill in the missing indices.
                  # Create padding first time we encounter data
                  if max(inds) < len(d):
                      # Simple case where we can simply pick slices from buffer
                      element_at_t = d[inds]
                  else:
                      # Case in which we have to pad because buffer has insufficient
                      # length. This branch takes more time than simply picking
                      # slices we try to avoid it.
                      element_at_t = _get_buffered_slice_with_paddings(
                          d, inds
                      )
                      element_at_t = np.stack(element_at_t)
  
                  if element_at_t.shape[0] == 1:
                      # Remove the T dimension if it is 1.
                      element_at_t = element_at_t[0]
                  shifted_data.append(element_at_t)
  
              # in some multi-agent cases shifted_data may be an empty list.
              # In this case we should just create an empty array and return it.
              if shifted_data:
                  shifted_data_np = np.stack(shifted_data, 0)
              else:
                  shifted_data_np = np.array(shifted_data)
              data.append(shifted_data_np)
  
          # We unflatten even if data is empty here, because the structure might be
          # nested with empty leafs and so we still need to reconstruct it.
          # This is useful because we spec-check states in RLModules and these
          # states can sometimes be nested dicts with empty leafs.
          batch_data[view_col] = self._unflatten_as_buffer_struct(
              data, data_col
          )
  
      batch = self._get_sample_batch(batch_data)
  
      # This trajectory is continuing -> Copy data at the end (in the size of
      # self.shift_before) to the beginning of buffers and erase everything
      # else.
      if (
          SampleBatch.TERMINATEDS in self.buffers
          and not self.buffers[SampleBatch.TERMINATEDS][0][-1]
          and SampleBatch.TRUNCATEDS in self.buffers
          and not self.buffers[SampleBatch.TRUNCATEDS][0][-1]
      ):
          # Copy data to beginning of buffer and cut lists.
          if self.shift_before > 0:
              for k, data in self.buffers.items():
                  # Loop through
                  for i in range(len(data)):
                      self.buffers[k][i] = data[i][-self.shift_before :]
          self.agent_steps = 0
  
      # Reset our unroll_id.
      self.unroll_id = None
  
      return batch
  
imports: import copy
  import logging
  import math
  from typing import Any, Dict, List, Optional
  import numpy
  import tree
  from gymnasium.spaces import Space
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray, get_dummy_batch_for_space
  from ray.rllib.utils.typing import EpisodeID, EnvID, TensorType, ViewRequirementsDict
  from ray.util.annotations import PublicAPI
label_desc: Use image node_id 8 for calling the AgentCollector obj's underlying member method code with example usage: obj.build_for_training(view_requirements) and returns: batch