class_docstr: Collect experiences from user environment using Connectors.
docstr: Process raw obs from env.
  
  Group data for active agents by policy. Reset environments that are done.
  
  Args:
      unfiltered_obs: The unfiltered, raw observations from the BaseEnv
          (vectorized, possibly multi-agent). Dict of dict: By env index,
          then agent ID, then mapped to actual obs.
      rewards: The rewards MultiEnvDict of the BaseEnv.
      terminateds: The `terminated` flags MultiEnvDict of the BaseEnv.
      truncateds: The `truncated` flags MultiEnvDict of the BaseEnv.
      infos: The MultiEnvDict of infos dicts of the BaseEnv.
  
  Returns:
      A tuple of:
          A list of envs that were active during this step.
          AgentConnectorDataType for active agents for policy evaluation.
          SampleBatches and RolloutMetrics for completed agents for output.
docstr_tok: ["Process","raw","obs","from","env",".","Group","data","for","active","agents","by","policy",".","Reset","environments","that","are","done",".","Args",":","unfiltered_obs",":","The","unfiltered",",","raw","observations","from","the","BaseEnv","(","vectorized",",","possibly","multi-agent",")",".","Dict","of","dict",":","By","env","index",",","then","agent","ID",",","then","mapped","to","actual","obs",".","rewards",":","The","rewards","MultiEnvDict","of","the","BaseEnv",".","terminateds",":","The","`","terminated","`","flags","MultiEnvDict","of","the","BaseEnv",".","truncateds",":","The","`","truncated","`","flags","MultiEnvDict","of","the","BaseEnv",".","infos",":","The","MultiEnvDict","of","infos","dicts","of","the","BaseEnv",".","Returns",":","A","tuple","of",":","A","list","of","envs","that","were","active","during","this","step",".","AgentConnectorDataType","for","active","agents","for","policy","evaluation",".","SampleBatches","and","RolloutMetrics","for","completed","agents","for","output","."]
code: def _process_observations(
      self,
      unfiltered_obs: MultiEnvDict,
      rewards: MultiEnvDict,
      terminateds: MultiEnvDict,
      truncateds: MultiEnvDict,
      infos: MultiEnvDict,
  ) -> Tuple[
      Set[EnvID],
      Dict[PolicyID, List[AgentConnectorDataType]],
      List[Union[RolloutMetrics, SampleBatchType]],
  ]:
      """Process raw obs from env.
  
      Group data for active agents by policy. Reset environments that are done.
  
      Args:
          unfiltered_obs: The unfiltered, raw observations from the BaseEnv
              (vectorized, possibly multi-agent). Dict of dict: By env index,
              then agent ID, then mapped to actual obs.
          rewards: The rewards MultiEnvDict of the BaseEnv.
          terminateds: The `terminated` flags MultiEnvDict of the BaseEnv.
          truncateds: The `truncated` flags MultiEnvDict of the BaseEnv.
          infos: The MultiEnvDict of infos dicts of the BaseEnv.
  
      Returns:
          A tuple of:
              A list of envs that were active during this step.
              AgentConnectorDataType for active agents for policy evaluation.
              SampleBatches and RolloutMetrics for completed agents for output.
      """
      # Output objects.
      # Note that we need to track envs that are active during this round explicitly,
      # just to be confident which envs require us to send at least an empty action
      # dict to.
      # We can not get this from the _active_episode or to_eval lists because
      # 1. All envs are not required to step during every single step. And
      # 2. to_eval only contains data for the agents that are still active. An env may
      # be active but all agents are done during the step.
      active_envs: Set[EnvID] = set()
      to_eval: Dict[
          PolicyID, List[AgentConnectorDataType]
      ] = defaultdict(list)
      outputs: List[Union[RolloutMetrics, SampleBatchType]] = []
  
      # For each (vectorized) sub-environment.
      # types: EnvID, Dict[AgentID, EnvObsType]
      for env_id, env_obs in unfiltered_obs.items():
          # Check for env_id having returned an error instead of a multi-agent
          # obs dict. This is how our BaseEnv can tell the caller to `poll()` that
          # one of its sub-environments is faulty and should be restarted (and the
          # ongoing episode should not be used for training).
          if isinstance(env_obs, Exception):
              assert terminateds[env_id]["__all__"] is True, (
                  f"ERROR: When a sub-environment (env-id {env_id}) returns an error "
                  "as observation, the terminateds[__all__] flag must also be set to "
                  "True!"
              )
              # all_agents_obs is an Exception here.
              # Drop this episode and skip to next.
              self._handle_done_episode(
                  env_id=env_id,
                  env_obs_or_exception=env_obs,
                  is_done=True,
                  active_envs=active_envs,
                  to_eval=to_eval,
                  outputs=outputs,
              )
              continue
  
          if env_id not in self._active_episodes:
              episode: EpisodeV2 = self.create_episode(env_id)
              self._active_episodes[env_id] = episode
          else:
              episode: EpisodeV2 = self._active_episodes[env_id]
          # If this episode is brand-new, call the episode start callback(s).
          # Note: EpisodeV2s are initialized with length=-1 (before the reset).
          if not episode.has_init_obs():
              self._call_on_episode_start(episode, env_id)
  
          # Check episode termination conditions.
          if (
              terminateds[env_id]["__all__"]
              or truncateds[env_id]["__all__"]
          ):
              all_agents_done = True
          else:
              all_agents_done = False
              active_envs.add(env_id)
  
          # Special handling of common info dict.
          episode.set_last_info(
              "__common__", infos[env_id].get("__common__", {})
          )
  
          # Agent sample batches grouped by policy. Each set of sample batches will
          # go through agent connectors together.
          sample_batches_by_policy = defaultdict(list)
          # Whether an agent is terminated or truncated.
          agent_terminateds = {}
          agent_truncateds = {}
          for agent_id, obs in env_obs.items():
              assert agent_id != "__all__"
  
              policy_id: PolicyID = episode.policy_for(agent_id)
  
              agent_terminated = bool(
                  terminateds[env_id]["__all__"]
                  or terminateds[env_id].get(agent_id)
              )
              agent_terminateds[agent_id] = agent_terminated
              agent_truncated = bool(
                  truncateds[env_id]["__all__"]
                  or truncateds[env_id].get(agent_id, False)
              )
              agent_truncateds[agent_id] = agent_truncated
  
              # A completely new agent is already done -> Skip entirely.
              if not episode.has_init_obs(agent_id) and (
                  agent_terminated or agent_truncated
              ):
                  continue
  
              values_dict = {
                  SampleBatch.T: episode.length,  # Episodes start at -1 before we
                  # add the initial obs. After that, we infer from initial obs at
                  # t=0 since that will be our new episode.length.
                  SampleBatch.ENV_ID: env_id,
                  SampleBatch.AGENT_INDEX: episode.agent_index(
                      agent_id
                  ),
                  # Last action (SampleBatch.ACTIONS) column will be populated by
                  # StateBufferConnector.
                  # Reward received after taking action at timestep t.
                  SampleBatch.REWARDS: rewards[env_id].get(
                      agent_id, 0.0
                  ),
                  # After taking action=a, did we reach terminal?
                  SampleBatch.TERMINATEDS: agent_terminated,
                  # Was the episode truncated artificially
                  # (e.g. b/c of some time limit)?
                  SampleBatch.TRUNCATEDS: agent_truncated,
                  SampleBatch.INFOS: infos[env_id].get(agent_id, {}),
                  SampleBatch.NEXT_OBS: obs,
              }
  
              # Queue this obs sample for connector preprocessing.
              sample_batches_by_policy[policy_id].append(
                  (agent_id, values_dict)
              )
  
          # The entire episode is done.
          if all_agents_done:
              # Let's check to see if there are any agents that haven't got the
              # last obs yet. If there are, we have to create fake-last
              # observations for them. (the environment is not required to do so if
              # terminateds[__all__]==True or truncateds[__all__]==True).
              for agent_id in episode.get_agents():
                  # If the latest obs we got for this agent is done, or if its
                  # episode state is already done, nothing to do.
                  if (
                      agent_terminateds.get(agent_id, False)
                      or agent_truncateds.get(agent_id, False)
                      or episode.is_done(agent_id)
                  ):
                      continue
  
                  policy_id: PolicyID = episode.policy_for(agent_id)
                  policy = self._worker.policy_map[policy_id]
  
                  # Create a fake observation by sampling the original env
                  # observation space.
                  obs_space = get_original_space(
                      policy.observation_space
                  )
                  # Although there is no obs for this agent, there may be
                  # good rewards and info dicts for it.
                  # This is the case for e.g. OpenSpiel games, where a reward
                  # is only earned with the last step, but the obs for that
                  # step is {}.
                  reward = rewards[env_id].get(agent_id, 0.0)
                  info = infos[env_id].get(agent_id, {})
                  values_dict = {
                      SampleBatch.T: episode.length,
                      SampleBatch.ENV_ID: env_id,
                      SampleBatch.AGENT_INDEX: episode.agent_index(
                          agent_id
                      ),
                      # TODO(sven): These should be the summed-up(!) rewards since the
                      #  last observation received for this agent.
                      SampleBatch.REWARDS: reward,
                      SampleBatch.TERMINATEDS: True,
                      SampleBatch.TRUNCATEDS: truncateds[env_id].get(
                          agent_id, False
                      ),
                      SampleBatch.INFOS: info,
                      SampleBatch.NEXT_OBS: obs_space.sample(),
                  }
  
                  # Queue these fake obs for connector preprocessing too.
                  sample_batches_by_policy[policy_id].append(
                      (agent_id, values_dict)
                  )
  
          # Run agent connectors.
          for policy_id, batches in sample_batches_by_policy.items():
              policy: Policy = self._worker.policy_map[policy_id]
              # Collected full MultiAgentDicts for this environment.
              # Run agent connectors.
              assert (
                  policy.agent_connectors
              ), "EnvRunnerV2 requires agent connectors to work."
  
              acd_list: List[AgentConnectorDataType] = [
                  AgentConnectorDataType(env_id, agent_id, data)
                  for agent_id, data in batches
              ]
  
              # For all agents mapped to policy_id, run their data
              # through agent_connectors.
              processed = policy.agent_connectors(acd_list)
  
              for d in processed:
                  # Record transition info if applicable.
                  if not episode.has_init_obs(d.agent_id):
                      episode.add_init_obs(
                          agent_id=d.agent_id,
                          init_obs=d.data.raw_dict[
                              SampleBatch.NEXT_OBS
                          ],
                          init_infos=d.data.raw_dict[SampleBatch.INFOS],
                          t=d.data.raw_dict[SampleBatch.T],
                      )
                  else:
                      episode.add_action_reward_done_next_obs(
                          d.agent_id, d.data.raw_dict
                      )
  
                  # Need to evaluate next actions.
                  if not (
                      all_agents_done
                      or agent_terminateds.get(d.agent_id, False)
                      or agent_truncateds.get(d.agent_id, False)
                      or episode.is_done(d.agent_id)
                  ):
                      # Add to eval set if env is not done and this particular agent
                      # is also not done.
                      item = AgentConnectorDataType(
                          d.env_id, d.agent_id, d.data
                      )
                      to_eval[policy_id].append(item)
  
          # Finished advancing episode by 1 step, mark it so.
          episode.step()
  
          # Exception: The very first env.poll() call causes the env to get reset
          # (no step taken yet, just a single starting observation logged).
          # We need to skip this callback in this case.
          if episode.length > 0:
              # Invoke the `on_episode_step` callback after the step is logged
              # to the episode.
              self._callbacks.on_episode_step(
                  worker=self._worker,
                  base_env=self._base_env,
                  policies=self._worker.policy_map,
                  episode=episode,
                  env_index=env_id,
              )
  
          # Episode is terminated/truncated for all agents
          # (terminateds[__all__] == True or truncateds[__all__] == True).
          if all_agents_done:
              # _handle_done_episode will build a MultiAgentBatch for all
              # the agents that are done during this step of rollout in
              # the case of _multiple_episodes_in_batch=False.
              self._handle_done_episode(
                  env_id,
                  env_obs,
                  terminateds[env_id]["__all__"]
                  or truncateds[env_id]["__all__"],
                  active_envs,
                  to_eval,
                  outputs,
              )
  
          # Try to build something.
          if self._multiple_episodes_in_batch:
              sample_batch = (
                  self._try_build_truncated_episode_multi_agent_batch(
                      self._batch_builders[env_id], episode
                  )
              )
              if sample_batch:
                  outputs.append(sample_batch)
  
                  # SampleBatch built from data collected by batch_builder.
                  # Clean up and delete the batch_builder.
                  del self._batch_builders[env_id]
  
      return active_envs, to_eval, outputs
  
imports: from collections import defaultdict
  import logging
  import time
  import tree
  from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Set, Tuple, Union
  import numpy
  from ray.rllib.env.base_env import ASYNC_RESET_RETURN, BaseEnv
  from ray.rllib.env.external_env import ExternalEnvWrapper
  from ray.rllib.env.wrappers.atari_wrappers import MonitorEnv, get_wrapper_by_cls
  from ray.rllib.evaluation.collectors.simple_list_collector import _PolicyCollectorGroup
  from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size
  from ray.rllib.evaluation.episode_v2 import EpisodeV2
  from ray.rllib.evaluation.metrics import RolloutMetrics
  from ray.rllib.models.preprocessors import Preprocessor
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, concat_samples
  from ray.rllib.utils.annotations import DeveloperAPI
  from ray.rllib.utils.filter import Filter
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.spaces.space_utils import unbatch, get_original_space
  from ray.rllib.utils.typing import ActionConnectorDataType, AgentConnectorDataType, AgentID, EnvActionType, EnvID, EnvInfoDict, EnvObsType, MultiAgentDict, MultiEnvDict, PolicyID, PolicyOutputType, SampleBatchType, StateBatches, TensorStructType
  from ray.util.debug import log_once
label_desc: Use image node_id 8 for calling the EnvRunnerV2 obj's underlying member method code with example usage: obj._process_observations(unfiltered_obs, rewards, terminateds, truncateds, infos) and returns: active_envs,  to_eval,  outputs