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 an initial observation (after reset) to the Agent's trajectory.
  
  Args:
      episode_id: Unique ID for the episode we are adding the
          initial observation for.
      agent_index: Unique int index (starting from 0) for the agent
          within its episode. Not to be confused with AGENT_ID (Any).
      env_id: The environment index (in a vectorized setup).
      init_obs: The initial observation tensor (after `env.reset()`).
      init_infos: The initial infos dict (after `env.reset()`).
      t: The time step (episode length - 1). The initial obs has
          ts=-1(!), then an action/reward/next-obs at t=0, etc..
docstr_tok: ["Adds","an","initial","observation","(","after","reset",")","to","the","Agent","'s","trajectory",".","Args",":","episode_id",":","Unique","ID","for","the","episode","we","are","adding","the","initial","observation","for",".","agent_index",":","Unique","int","index","(","starting","from","0",")","for","the","agent","within","its","episode",".","Not","to","be","confused","with","AGENT_ID","(","Any",")",".","env_id",":","The","environment","index","(","in","a","vectorized","setup",")",".","init_obs",":","The","initial","observation","tensor","(","after","`","env.reset","(",")","`",")",".","init_infos",":","The","initial","infos","dict","(","after","`","env.reset","(",")","`",")",".","t",":","The","time","step","(","episode","length","-","1",")",".","The","initial","obs","has","ts=-1","(","!",")",",","then","an","action\/reward\/next-obs","at","t=0",",","etc",".."]
code: def add_init_obs(
      self,
      episode_id: EpisodeID,
      agent_index: int,
      env_id: EnvID,
      init_obs: TensorType,
      init_infos: Optional[Dict[str, TensorType]] = None,
      t: int = -1,
  ) -> None:
      """Adds an initial observation (after reset) to the Agent's trajectory.
  
      Args:
          episode_id: Unique ID for the episode we are adding the
              initial observation for.
          agent_index: Unique int index (starting from 0) for the agent
              within its episode. Not to be confused with AGENT_ID (Any).
          env_id: The environment index (in a vectorized setup).
          init_obs: The initial observation tensor (after `env.reset()`).
          init_infos: The initial infos dict (after `env.reset()`).
          t: The time step (episode length - 1). The initial obs has
              ts=-1(!), then an action/reward/next-obs at t=0, etc..
      """
      # Store episode ID + unroll ID, which will be constant throughout this
      # AgentCollector's lifecycle.
      self.episode_id = episode_id
      if self.unroll_id is None:
          self.unroll_id = AgentCollector._next_unroll_id
          AgentCollector._next_unroll_id += 1
  
      # convert init_obs to np.array (in case it is a list)
      if isinstance(init_obs, list):
          init_obs = np.array(init_obs)
  
      if SampleBatch.OBS not in self.buffers:
          single_row = {
              SampleBatch.OBS: init_obs,
              SampleBatch.INFOS: init_infos or {},
              SampleBatch.AGENT_INDEX: agent_index,
              SampleBatch.ENV_ID: env_id,
              SampleBatch.T: t,
              SampleBatch.EPS_ID: self.episode_id,
              SampleBatch.UNROLL_ID: self.unroll_id,
          }
  
          # TODO (Artur): Remove when PREV_ACTIONS and PREV_REWARDS get deprecated.
          # Note (Artur): As long as we have these in our default view requirements,
          # we should  build buffers with neutral elements instead of building them
          # on the first AgentCollector.build_for_inference call if present.
          # This prevents us from accidentally building buffers with duplicates of
          # the first incoming value.
          if SampleBatch.PREV_REWARDS in self.view_requirements:
              single_row[
                  SampleBatch.REWARDS
              ] = get_dummy_batch_for_space(
                  space=self.view_requirements[
                      SampleBatch.REWARDS
                  ].space,
                  batch_size=0,
                  fill_value=0.0,
              )
          if SampleBatch.PREV_ACTIONS in self.view_requirements:
              potentially_flattened_batch = get_dummy_batch_for_space(
                  space=self.view_requirements[
                      SampleBatch.ACTIONS
                  ].space,
                  batch_size=0,
                  fill_value=0.0,
              )
              if not self.disable_action_flattening:
                  potentially_flattened_batch = (
                      flatten_to_single_ndarray(
                          potentially_flattened_batch
                      )
                  )
              single_row[
                  SampleBatch.ACTIONS
              ] = potentially_flattened_batch
          self._build_buffers(single_row)
  
      # Append data to existing buffers.
      flattened = tree.flatten(init_obs)
      for i, sub_obs in enumerate(flattened):
          self.buffers[SampleBatch.OBS][i].append(sub_obs)
      self.buffers[SampleBatch.INFOS][0].append(init_infos or {})
      self.buffers[SampleBatch.AGENT_INDEX][0].append(agent_index)
      self.buffers[SampleBatch.ENV_ID][0].append(env_id)
      self.buffers[SampleBatch.T][0].append(t)
      self.buffers[SampleBatch.EPS_ID][0].append(self.episode_id)
      self.buffers[SampleBatch.UNROLL_ID][0].append(self.unroll_id)
  
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 5 for calling the AgentCollector obj's underlying member method code with example usage: obj.add_init_obs(episode_id, agent_index, env_id, init_obs, init_infos, t) without return types