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: Initialize an AgentCollector.
  
  Args:
      view_reqs: A dict of view requirements for the agent.
      max_seq_len: The maximum sequence length to store.
      disable_action_flattening: If True, don't flatten the action.
      intial_states: The initial states from the policy.get_initial_states()
      is_policy_recurrent: If True, the policy is recurrent.
      is_training: Sets the is_training flag for the buffers. if True, all the
          timesteps are stored in the buffers until explictly build_for_training
          () is called. if False, only the content required for the last time
          step is stored in the buffers. This will save memory during inference.
          You can change the behavior at runtime by calling is_training(mode).
docstr_tok: ["Initialize","an","AgentCollector",".","Args",":","view_reqs",":","A","dict","of","view","requirements","for","the","agent",".","max_seq_len",":","The","maximum","sequence","length","to","store",".","disable_action_flattening",":","If","True",",","do","n't","flatten","the","action",".","intial_states",":","The","initial","states","from","the","policy.get_initial_states","(",")","is_policy_recurrent",":","If","True",",","the","policy","is","recurrent",".","is_training",":","Sets","the","is_training","flag","for","the","buffers",".","if","True",",","all","the","timesteps","are","stored","in","the","buffers","until","explictly","build_for_training","(",")","is","called",".","if","False",",","only","the","content","required","for","the","last","time","step","is","stored","in","the","buffers",".","This","will","save","memory","during","inference",".","You","can","change","the","behavior","at","runtime","by","calling","is_training","(","mode",")","."]
code: def __init__(
      self,
      view_reqs: ViewRequirementsDict,
      *,
      max_seq_len: int = 1,
      disable_action_flattening: bool = True,
      intial_states: Optional[List[TensorType]] = None,
      is_policy_recurrent: bool = False,
      is_training: bool = True,
      _enable_new_api_stack: bool = False,
  ):
      """Initialize an AgentCollector.
  
      Args:
          view_reqs: A dict of view requirements for the agent.
          max_seq_len: The maximum sequence length to store.
          disable_action_flattening: If True, don't flatten the action.
          intial_states: The initial states from the policy.get_initial_states()
          is_policy_recurrent: If True, the policy is recurrent.
          is_training: Sets the is_training flag for the buffers. if True, all the
              timesteps are stored in the buffers until explictly build_for_training
              () is called. if False, only the content required for the last time
              step is stored in the buffers. This will save memory during inference.
              You can change the behavior at runtime by calling is_training(mode).
      """
      self.max_seq_len = max_seq_len
      self.disable_action_flattening = disable_action_flattening
      self.view_requirements = view_reqs
      # The initial_states can be an np array
      self.initial_states = (
          intial_states if intial_states is not None else []
      )
      self.is_policy_recurrent = is_policy_recurrent
      self._is_training = is_training
      self._enable_new_api_stack = _enable_new_api_stack
  
      # Determine the size of the buffer we need for data before the actual
      # episode starts. This is used for 0-buffering of e.g. prev-actions,
      # or internal state inputs.
      view_req_shifts = [
          min(vr.shift_arr)
          - int(
              (vr.data_col or k) in [SampleBatch.OBS, SampleBatch.INFOS]
          )
          for k, vr in view_reqs.items()
      ]
      self.shift_before = -min(view_req_shifts)
  
      # The actual data buffers. Keys are column names, values are lists
      # that contain the sub-components (e.g. for complex obs spaces) with
      # each sub-component holding a list of per-timestep tensors.
      # E.g.: obs-space = Dict(a=Discrete(2), b=Box((2,)))
      # buffers["obs"] = [
      #    [0, 1],  # <- 1st sub-component of observation
      #    [np.array([.2, .3]), np.array([.0, -.2])]  # <- 2nd sub-component
      # ]
      # NOTE: infos and state_out... are not flattened due to them often
      # using custom dict values whose structure may vary from timestep to
      # timestep.
      self.buffers: Dict[str, List[List[TensorType]]] = {}
      # Maps column names to an example data item, which may be deeply
      # nested. These are used such that we'll know how to unflatten
      # the flattened data inside self.buffers when building the
      # SampleBatch.
      self.buffer_structs: Dict[str, Any] = {}
      # The episode ID for the agent for which we collect data.
      self.episode_id = None
      # The unroll ID, unique across all rollouts (within a RolloutWorker).
      self.unroll_id = None
      # The simple timestep count for this agent. Gets increased by one
      # each time a (non-initial!) observation is added.
      self.agent_steps = 0
      # Keep track of view requirements that have a view on columns that we gain from
      # inference and also need for inference. These have dummy values appended in
      # buffers to account for the missing value when building for inference
      # Example: We have one 'state_in' view requirement that has a view on our
      # state_outs at t=[-10, ..., -1]. At any given build_for_inference()-call,
      # the buffer must contain eleven values from t=[-10, ..., 0] for us to index
      # properly. Since state_out at t=0 is missing, we substitute it with a buffer
      # value that should never make it into batches built for training.
      self.data_cols_with_dummy_values = set()
  
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 1 to create a new AgentCollector object  with example: obj = AgentCollector(view_reqs)