code: def _process_observations(
      *,
      worker: "RolloutWorker",
      base_env: BaseEnv,
      active_episodes: Dict[EnvID, Episode],
      unfiltered_obs: Dict[EnvID, Dict[AgentID, EnvObsType]],
      rewards: Dict[EnvID, Dict[AgentID, float]],
      terminateds: Dict[EnvID, Dict[AgentID, bool]],
      truncateds: Dict[EnvID, Dict[AgentID, bool]],
      infos: Dict[EnvID, Dict[AgentID, EnvInfoDict]],
      multiple_episodes_in_batch: bool,
      callbacks: "DefaultCallbacks",
      observation_fn: "ObservationFunction",
      sample_collector: SampleCollector,
  ) -> Tuple[
      Set[EnvID],
      Dict[PolicyID, List[_PolicyEvalData]],
      List[Union[RolloutMetrics, SampleBatchType]],
  ]:
      """Record new data from the environment and prepare for policy evaluation.
  
      Args:
          worker: Reference to the current rollout worker.
          base_env: Env implementing BaseEnv.
          active_episodes: Mapping from
              episode ID to currently ongoing Episode object.
          unfiltered_obs: Doubly keyed dict of env-ids -> agent ids
              -> unfiltered observation tensor, returned by a `BaseEnv.poll()`
              call.
          rewards: Doubly keyed dict of env-ids -> agent ids ->
              rewards tensor, returned by a `BaseEnv.poll()` call.
          terminateds: Doubly keyed dict of env-ids -> agent ids ->
              boolean `terminated` flags, returned by a `BaseEnv.poll()` call.
          truncateds: Doubly keyed dict of env-ids -> agent ids ->
              boolean `truncated` flags, returned by a `BaseEnv.poll()` call.
          infos: Doubly keyed dict of env-ids -> agent ids ->
              info dicts, returned by a `BaseEnv.poll()` call.
          multiple_episodes_in_batch: Whether to pack multiple
              episodes into each batch. This guarantees batches will be exactly
              `rollout_fragment_length` in size.
          callbacks: User callbacks to run on episode events.
          observation_fn: Optional multi-agent
              observation func to use for preprocessing observations.
          sample_collector: The SampleCollector object
              used to store and retrieve environment samples.
  
      Returns:
          Tuple consisting of 1) active_envs: Set of non-terminated env ids.
          2) to_eval: Map of policy_id to list of agent _PolicyEvalData.
          3) outputs: List of metrics and samples to return from the sampler.
      """
  
      # Output objects.
      active_envs: Set[EnvID] = set()
      to_eval: Dict[PolicyID, List[_PolicyEvalData]] = defaultdict(list)
      outputs: List[Union[RolloutMetrics, SampleBatchType]] = []
  
      # For each (vectorized) sub-environment.
      # types: EnvID, Dict[AgentID, EnvObsType]
      for env_id, all_agents_obs in unfiltered_obs.items():
          episode: Episode = active_episodes[env_id]
  
          # 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(all_agents_obs, Exception):
              episode.is_faulty = True
              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!"
              )
              # This will be filled with dummy observations below.
              all_agents_obs = {}
  
          # Add init obs and infos (from the call to `reset/try_reset`) to episode.
          for aid, obs in all_agents_obs.items():
              episode._set_last_raw_obs(aid, obs)
          common_infos = infos[env_id].get("__common__", {})
          episode._set_last_info("__common__", common_infos)
          for aid, info in infos[env_id].items():
              episode._set_last_info(aid, info)
  
          # Episode is brand new.
          if episode.started is False:
              # Call the episode start callback(s).
              _call_on_episode_start(
                  episode, env_id, callbacks, worker, base_env
              )
          else:
              sample_collector.episode_step(episode)
              episode._add_agent_rewards(rewards[env_id])
  
          # Check episode termination conditions.
          if (
              terminateds[env_id]["__all__"]
              or truncateds[env_id]["__all__"]
          ):
              all_agents_done = True
              # Check whether we have to create a fake-last observation
              # for some agents (the environment is not required to do so if
              # terminateds[__all__]=True or truncateds[__all__]=True).
              for ag_id in episode.get_agents():
                  if (
                      not episode.last_terminated_for(ag_id)
                      and not episode.last_truncated_for(ag_id)
                      and ag_id not in all_agents_obs
                  ):
                      # Create a fake (all-0s) observation.
                      obs_sp = worker.policy_map[
                          episode.policy_for(ag_id)
                      ].observation_space
                      obs_sp = getattr(obs_sp, "original_space", obs_sp)
                      all_agents_obs[ag_id] = tree.map_structure(
                          np.zeros_like, obs_sp.sample()
                      )
          else:
              all_agents_done = False
              active_envs.add(env_id)
  
          # Custom observation function is applied before preprocessing.
          if observation_fn:
              all_agents_obs: Dict[
                  AgentID, EnvObsType
              ] = observation_fn(
                  agent_obs=all_agents_obs,
                  worker=worker,
                  base_env=base_env,
                  policies=worker.policy_map,
                  episode=episode,
              )
              if not isinstance(all_agents_obs, dict):
                  raise ValueError(
                      "observe() must return a dict of agent observations"
                  )
  
          # For each agent in the environment.
          # types: AgentID, EnvObsType
          for agent_id, raw_obs in all_agents_obs.items():
              assert agent_id != "__all__"
  
              last_observation: EnvObsType = (
                  episode.last_observation_for(agent_id)
              )
              agent_terminated = bool(
                  terminateds[env_id]["__all__"]
                  or terminateds[env_id].get(agent_id)
              )
              agent_truncated = bool(
                  truncateds[env_id]["__all__"]
                  or truncateds[env_id].get(agent_id, False)
              )
  
              # A new agent (initial obs) is already done -> Skip entirely.
              if last_observation is None and (
                  agent_terminated or agent_truncated
              ):
                  continue
  
              policy_id: PolicyID = episode.policy_for(agent_id)
  
              preprocessor = _get_or_raise(
                  worker.preprocessors, policy_id
              )
              prep_obs: EnvObsType = raw_obs
              if preprocessor is not None:
                  prep_obs = preprocessor.transform(raw_obs)
                  if log_once("prep_obs"):
                      logger.info(
                          "Preprocessed obs: {}".format(
                              summarize(prep_obs)
                          )
                      )
              filtered_obs: EnvObsType = _get_or_raise(
                  worker.filters, policy_id
              )(prep_obs)
              if log_once("filtered_obs"):
                  logger.info(
                      "Filtered obs: {}".format(summarize(filtered_obs))
                  )
  
              episode._set_last_observation(agent_id, filtered_obs)
              episode._set_last_terminated(agent_id, agent_terminated)
              episode._set_last_truncated(agent_id, agent_truncated)
              agent_infos = infos[env_id].get(agent_id, {})
  
              # Record transition info if applicable.
              if last_observation is None:
                  sample_collector.add_init_obs(
                      episode=episode,
                      agent_id=agent_id,
                      env_id=env_id,
                      policy_id=policy_id,
                      init_obs=filtered_obs,
                      init_infos=agent_infos,
                      t=episode.length - 1,
                  )
              else:
                  # Add actions, rewards, next-obs to collectors.
                  values_dict = {
                      SampleBatch.T: episode.length - 1,
                      SampleBatch.ENV_ID: env_id,
                      SampleBatch.AGENT_INDEX: episode._agent_index(
                          agent_id
                      ),
                      # Action (slot 0) taken at timestep t.
                      SampleBatch.ACTIONS: episode.last_action_for(
                          agent_id
                      ),
                      # Reward received after taking a at timestep t.
                      SampleBatch.REWARDS: rewards[env_id].get(
                          agent_id, 0.0
                      ),
                      # After taking action=a, did we terminate the episode?
                      SampleBatch.TERMINATEDS: agent_terminated,
                      # Was the episode truncated artificially
                      # (e.g. b/c of some time limit)?
                      SampleBatch.TRUNCATEDS: agent_truncated,
                      # Next observation.
                      SampleBatch.NEXT_OBS: filtered_obs,
                  }
                  # Add extra-action-fetches (policy-inference infos) to
                  # collectors.
                  pol = worker.policy_map[policy_id]
                  for key, value in episode.last_extra_action_outs_for(
                      agent_id
                  ).items():
                      if key in pol.view_requirements:
                          values_dict[key] = value
                  # Env infos for this agent.
                  if SampleBatch.INFOS in pol.view_requirements:
                      values_dict[SampleBatch.INFOS] = agent_infos
                  sample_collector.add_action_reward_next_obs(
                      episode.episode_id,
                      agent_id,
                      env_id,
                      policy_id,
                      agent_terminated or agent_truncated,
                      values_dict,
                  )
  
              if not agent_terminated and not agent_truncated:
                  item = _PolicyEvalData(
                      env_id,
                      agent_id,
                      filtered_obs,
                      agent_infos,
                      None
                      if last_observation is None
                      else episode.rnn_state_for(agent_id),
                      None
                      if last_observation is None
                      else episode.last_action_for(agent_id),
                      rewards[env_id].get(agent_id, 0.0),
                  )
                  to_eval[policy_id].append(item)
  
          # Invoke the `on_episode_step` callback after the step is logged
          # to the episode.
          # 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 not episode.is_faulty and episode.length > 0:
              callbacks.on_episode_step(
                  worker=worker,
                  base_env=base_env,
                  policies=worker.policy_map,
                  episode=episode,
                  env_index=env_id,
              )
  
          # Episode is terminated for all agents (terminateds[__all__] == True or
          # truncateds[__all__] == True).
          if all_agents_done:
              # If, we are not allowed to pack the next episode into the same
              # SampleBatch (batch_mode=complete_episodes) -> Build the
              # MultiAgentBatch from a single episode and add it to "outputs".
              # Otherwise, just postprocess and continue collecting across
              # episodes.
              # If an episode was marked faulty, perform regular postprocessing
              # (to e.g. properly flush and clean up the SampleCollector's buffers),
              # but then discard the entire batch and don't return it.
              ma_sample_batch = None
              if not episode.is_faulty or episode.length > 0:
                  ma_sample_batch = (
                      sample_collector.postprocess_episode(
                          episode,
                          is_done=True,
                          check_dones=True,
                          build=episode.is_faulty
                          or not multiple_episodes_in_batch,
                      )
                  )
              if not episode.is_faulty:
                  # Call each (in-memory) policy's Exploration.on_episode_end
                  # method.
                  # Note: This may break the exploration (e.g. ParameterNoise) of
                  # policies in the `policy_map` that have not been recently used
                  # (and are therefore stashed to disk). However, we certainly do not
                  # want to loop through all (even stashed) policies here as that
                  # would counter the purpose of the LRU policy caching.
                  for p in worker.policy_map.cache.values():
                      if getattr(p, "exploration", None) is not None:
                          p.exploration.on_episode_end(
                              policy=p,
                              environment=base_env,
                              episode=episode,
                              tf_sess=p.get_session(),
                          )
                  # Call custom on_episode_end callback.
                  callbacks.on_episode_end(
                      worker=worker,
                      base_env=base_env,
                      policies=worker.policy_map,
                      episode=episode,
                      env_index=env_id,
                  )
  
              # Now that all callbacks are done and users had the chance to add custom
              # metrics based on the last observation in the episode, finish up metrics
              # object and append to `outputs`.
              atari_metrics: List[
                  RolloutMetrics
              ] = _fetch_atari_metrics(base_env)
              if not episode.is_faulty:
                  if atari_metrics is not None:
                      for m in atari_metrics:
                          outputs.append(
                              m._replace(
                                  custom_metrics=episode.custom_metrics,
                                  hist_data=episode.hist_data,
                              )
                          )
                  else:
                      outputs.append(
                          RolloutMetrics(
                              episode.length,
                              episode.total_reward,
                              dict(episode.agent_rewards),
                              episode.custom_metrics,
                              {},
                              episode.hist_data,
                              episode.media,
                          )
                      )
              else:
                  # Add metrics about a faulty episode.
                  outputs.append(RolloutMetrics(episode_faulty=True))
  
              # Only after the RolloutMetrics were appended, append the collected sample
              # batch, if any.
              if not episode.is_faulty and ma_sample_batch:
                  outputs.append(ma_sample_batch)
  
              # Terminated: Try to reset the sub environment.
              # Clean up old finished episode.
              del active_episodes[env_id]
  
              # Create a new episode and call `on_episode_created` callback(s).
              _create_episode(
                  active_episodes, env_id, callbacks, worker, base_env
              )
  
              # The sub environment at index `env_id` might throw an exception
              # during the following `try_reset()` attempt. If configured with
              # `restart_failed_sub_environments=True`, the BaseEnv will restart
              # the affected sub environment (create a new one using its c'tor) and
              # must reset the recreated sub env right after that.
              # Should the sub environment fail indefinitely during these
              # repeated reset attempts, the entire worker will be blocked.
              # This would be ok, b/c the alternative would be the worker crashing
              # entirely.
              while True:
                  resetted_obs, resetted_infos = base_env.try_reset(
                      env_id
                  )
                  if resetted_obs is None or not isinstance(
                      resetted_obs[env_id], Exception
                  ):
                      break
                  else:
                      # Failed to reset, add metrics about a faulty episode.
                      outputs.append(
                          RolloutMetrics(episode_faulty=True)
                      )
  
              # Creates a new episode if this is not async return.
              # If reset is async, we will get its result in some future poll.
              if (
                  resetted_obs is not None
                  and resetted_obs != ASYNC_RESET_RETURN
              ):
                  new_episode: Episode = active_episodes[env_id]
  
                  resetted_obs = resetted_obs[env_id]
                  resetted_infos = resetted_infos[env_id]
  
                  # Add init obs and infos (from the call to `reset/try_reset`) to
                  # episode.
                  for aid, obs in resetted_obs.items():
                      new_episode._set_last_raw_obs(aid, obs)
                  common_infos = resetted_infos.get("__common__", {})
                  new_episode._set_last_info("__common__", common_infos)
                  for aid, info in resetted_infos.items():
                      new_episode._set_last_info(aid, info)
  
                  _call_on_episode_start(
                      new_episode, env_id, callbacks, worker, base_env
                  )
  
                  _assert_episode_not_faulty(new_episode)
                  if observation_fn:
                      resetted_obs: Dict[
                          AgentID, EnvObsType
                      ] = observation_fn(
                          agent_obs=resetted_obs,
                          worker=worker,
                          base_env=base_env,
                          policies=worker.policy_map,
                          episode=new_episode,
                      )
                  # types: AgentID, EnvObsType
                  for agent_id, raw_obs in resetted_obs.items():
                      policy_id: PolicyID = new_episode.policy_for(
                          agent_id
                      )
                      preproccessor = _get_or_raise(
                          worker.preprocessors, policy_id
                      )
  
                      prep_obs: EnvObsType = raw_obs
                      if preproccessor is not None:
                          prep_obs = preproccessor.transform(raw_obs)
                      filtered_obs: EnvObsType = _get_or_raise(
                          worker.filters, policy_id
                      )(prep_obs)
                      new_episode._set_last_observation(
                          agent_id, filtered_obs
                      )
  
                      # Add initial obs to buffer.
                      sample_collector.add_init_obs(
                          episode=new_episode,
                          agent_id=agent_id,
                          env_id=env_id,
                          policy_id=policy_id,
                          init_obs=filtered_obs,
                          init_infos=resetted_infos,
                          t=new_episode.length - 1,
                      )
  
                      item = _PolicyEvalData(
                          env_id,
                          agent_id,
                          filtered_obs,
                          new_episode.last_info_for(agent_id) or {},
                          new_episode.rnn_state_for(agent_id),
                          None,
                          0.0,
                      )
                      to_eval[policy_id].append(item)
  
      # Try to build something.
      if multiple_episodes_in_batch:
          sample_batches = (
              sample_collector.try_build_truncated_episode_multi_agent_batch()
          )
          if sample_batches:
              outputs.extend(sample_batches)
  
      return active_envs, to_eval, outputs
  
imports: import logging
  import queue
  import time
  from abc import ABCMeta, abstractmethod
  from collections import defaultdict, namedtuple
  from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union
  import numpy
  import tree
  from ray.rllib.env.base_env import ASYNC_RESET_RETURN, BaseEnv, convert_to_base_env
  from ray.rllib.evaluation.collectors.sample_collector import SampleCollector
  from ray.rllib.evaluation.collectors.simple_list_collector import SimpleListCollector
  from ray.rllib.evaluation.env_runner_v2 import EnvRunnerV2, _fetch_atari_metrics, _get_or_raise, _PerfStats
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.metrics import RolloutMetrics
  from ray.rllib.offline import InputReader
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.policy_map import PolicyMap
  from ray.rllib.policy.sample_batch import SampleBatch, concat_samples
  from ray.rllib.utils.annotations import DeveloperAPI, override
  from ray.rllib.utils.debug import summarize
  from ray.rllib.utils.deprecation import deprecation_warning, DEPRECATED_VALUE
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.numpy import convert_to_numpy, make_action_immutable
  from ray.rllib.utils.spaces.space_utils import clip_action, unbatch, unsquash_action
  from ray.rllib.utils.typing import AgentID, EnvActionType, EnvID, EnvInfoDict, EnvObsType, MultiEnvDict, PolicyID, SampleBatchType, TensorStructType
  from ray.util.debug import log_once
label_desc: Use image node_id 2 for calling a global function with example usage: _process_observations() and returns: active_envs,  to_eval,  outputs