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: During inference, we will build a SampleBatch with a batch size of 1 that
  can then be used to run the forward pass of a policy. This data will only
  include the enviornment context for running the policy at the last timestep.
  
  Returns:
      A SampleBatch with a batch size of 1.
docstr_tok: ["During","inference",",","we","will","build","a","SampleBatch","with","a","batch","size","of","1","that","can","then","be","used","to","run","the","forward","pass","of","a","policy",".","This","data","will","only","include","the","enviornment","context","for","running","the","policy","at","the","last","timestep",".","Returns",":","A","SampleBatch","with","a","batch","size","of","1","."]
code: def build_for_inference(self) -> SampleBatch:
      """During inference, we will build a SampleBatch with a batch size of 1 that
      can then be used to run the forward pass of a policy. This data will only
      include the enviornment context for running the policy at the last timestep.
  
      Returns:
          A SampleBatch with a batch size of 1.
      """
  
      batch_data = {}
      np_data = {}
      for view_col, view_req in self.view_requirements.items():
          # Create the batch of data from the different buffers.
          data_col = view_req.data_col or view_col
  
          # if this view is not for inference, skip it.
          if not view_req.used_for_compute_actions:
              continue
  
          if np.any(view_req.shift_arr > 0):
              raise ValueError(
                  f"During inference the agent can only use past observations to "
                  f"respect causality. However, view_col = {view_col} seems to "
                  f"depend on future indices {view_req.shift_arr}, while the "
                  f"used_for_compute_actions flag is set to True. Please fix the "
                  f"discrepancy. Hint: If you are using a custom model make sure "
                  f"the view_requirements are initialized properly and is point "
                  f"only refering to past timesteps during inference."
              )
  
          # Some columns don't exist yet
          # (get created during postprocessing or depend on state_out).
          if data_col not in self.buffers:
              self._fill_buffer_with_initial_values(
                  data_col, view_req, build_for_inference=True
              )
              self._prepare_for_data_cols_with_dummy_values(data_col)
  
          # 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)
  
          data = []
          for d in np_data[data_col]:
              # if shift_arr = [0] the data will be just the last time step
              # (len(d) - 1), if shift_arr = [-1] the data will be just the timestep
              # before the last one (len(d) - 2) and so on.
              element_at_t = d[view_req.shift_arr + len(d) - 1]
              if element_at_t.shape[0] == 1:
                  # We'd normally squeeze here to remove the time dim, but we'll
                  # simply use the time dim as the batch dim.
                  data.append(element_at_t)
                  continue
              # add the batch dimension with [None]
              data.append(element_at_t[None])
  
          # 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)
      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 7 for calling the AgentCollector obj's underlying member method code with example usage: obj.build_for_inference() and returns: batch