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: Adds the given dictionary (row) of values to the Agent's trajectory.
  
  Args:
      values: Data dict (interpreted as a single row) to be added to buffer.
          Must contain keys:
          SampleBatch.ACTIONS, REWARDS, TERMINATEDS, TRUNCATEDS, and NEXT_OBS.
docstr_tok: ["Adds","the","given","dictionary","(","row",")","of","values","to","the","Agent","'s","trajectory",".","Args",":","values",":","Data","dict","(","interpreted","as","a","single","row",")","to","be","added","to","buffer",".","Must","contain","keys",":","SampleBatch.ACTIONS",",","REWARDS",",","TERMINATEDS",",","TRUNCATEDS",",","and","NEXT_OBS","."]
code: def add_action_reward_next_obs(
      self, input_values: Dict[str, TensorType]
  ) -> None:
      """Adds the given dictionary (row) of values to the Agent's trajectory.
  
      Args:
          values: Data dict (interpreted as a single row) to be added to buffer.
              Must contain keys:
              SampleBatch.ACTIONS, REWARDS, TERMINATEDS, TRUNCATEDS, and NEXT_OBS.
      """
      if self.unroll_id is None:
          self.unroll_id = AgentCollector._next_unroll_id
          AgentCollector._next_unroll_id += 1
  
      # Next obs -> obs.
      values = copy.copy(input_values)
      assert SampleBatch.OBS not in values
      values[SampleBatch.OBS] = values[SampleBatch.NEXT_OBS]
      del values[SampleBatch.NEXT_OBS]
  
      # convert obs to np.array (in case it is a list)
      if isinstance(values[SampleBatch.OBS], list):
          values[SampleBatch.OBS] = np.array(values[SampleBatch.OBS])
  
      # Default to next timestep if not provided in input values
      if SampleBatch.T not in input_values:
          values[SampleBatch.T] = self.buffers[SampleBatch.T][0][-1] + 1
  
      # Make sure EPS_ID/UNROLL_ID stay the same for this agent.
      if SampleBatch.EPS_ID in values:
          assert values[SampleBatch.EPS_ID] == self.episode_id
          del values[SampleBatch.EPS_ID]
      self.buffers[SampleBatch.EPS_ID][0].append(self.episode_id)
      if SampleBatch.UNROLL_ID in values:
          assert values[SampleBatch.UNROLL_ID] == self.unroll_id
          del values[SampleBatch.UNROLL_ID]
      self.buffers[SampleBatch.UNROLL_ID][0].append(self.unroll_id)
  
      for k, v in values.items():
          if k not in self.buffers:
              if self.training and k.startswith("state_out"):
                  vr = self.view_requirements[k]
                  data_col = vr.data_col or k
                  self._fill_buffer_with_initial_values(
                      data_col, vr, build_for_inference=False
                  )
              else:
                  self._build_buffers({k: v})
          # Do not flatten infos, state_out and (if configured) actions.
          # Infos/state-outs may be structs that change from timestep to
          # timestep.
          should_flatten_action_key = (
              k == SampleBatch.ACTIONS
              and not self.disable_action_flattening
          )
          # Note (Artur) RL Modules's states need no flattening
          should_flatten_state_key = (
              k.startswith("state_out")
              and not self._enable_new_api_stack
          )
          if (
              k == SampleBatch.INFOS
              or should_flatten_state_key
              or should_flatten_action_key
          ):
              if should_flatten_action_key:
                  v = flatten_to_single_ndarray(v)
              # Briefly remove dummy value to add to buffer
              if k in self.data_cols_with_dummy_values:
                  dummy = self.buffers[k][0].pop(-1)
              self.buffers[k][0].append(v)
              # Add back dummy value
              if k in self.data_cols_with_dummy_values:
                  self.buffers[k][0].append(dummy)
          # Flatten all other columns.
          else:
              flattened = tree.flatten(v)
              for i, sub_list in enumerate(self.buffers[k]):
                  # Briefly remove dummy value to add to buffer
                  if k in self.data_cols_with_dummy_values:
                      dummy = sub_list.pop(-1)
                  sub_list.append(flattened[i])
                  # Add back dummy value
                  if k in self.data_cols_with_dummy_values:
                      sub_list.append(dummy)
  
      # In inference mode, we don't need to keep all of trajectory in memory
      # we only need to keep the steps required. We can pop from the beginning to
      # create room for new data.
      if not self.training:
          for k in self.buffers:
              for sub_list in self.buffers[k]:
                  if sub_list:
                      sub_list.pop(0)
  
      self.agent_steps += 1
  
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 6 for calling the AgentCollector obj's underlying member method code with example usage: obj.add_action_reward_next_obs(input_values) without return types