class_docstr: Stores multi-agent episode data.
  
  The central attribute of the class is the timestep mapping
  `global_t_to_local_t` that maps the global (environment)
  timestep to the local (agent) timesteps.
  
  The `MultiAgentEpisode` is based on the `SingleAgentEpisode`s
  for each agent, stored in `MultiAgentEpisode.agent_episodes`.
class_docstr_tok: ["Stores","multi-agent","episode","data",".","The","central","attribute","of","the","class","is","the","timestep","mapping","`","global_t_to_local_t","`","that","maps","the","global","(","environment",")","timestep","to","the","local","(","agent",")","timesteps",".","The","`","MultiAgentEpisode","`","is","based","on","the","`","SingleAgentEpisode","`","s","for","each","agent",",","stored","in","`","MultiAgentEpisode.agent_episodes","`","."]
docstr: Adds a timestep to the episode.
  
  Args:
      observations: A dictionary mapping agent ids to their corresponding
          observations. Note that some agents may not have stepped at this
          timestep.
      actions: Mandatory. A dictionary mapping agent ids to their
          corresponding actions. Note that some agents may not have stepped at
          this timestep.
      rewards: Mandatory. A dictionary mapping agent ids to their
          corresponding observations. Note that some agents may not have stepped
          at this timestep.
      infos: A dictionary mapping agent ids to their
          corresponding info. Note that some agents may not have stepped at this
          timestep.
      terminateds: A dictionary mapping agent ids to their `terminated` flags,
          indicating, whether the environment has been terminated for them.
          A special `__all__` key indicates that the episode is terminated for
          all agent ids.
      terminateds: A dictionary mapping agent ids to their `truncated` flags,
          indicating, whether the environment has been truncated for them.
          A special `__all__` key indicates that the episode is `truncated` for
          all agent ids.
      render_image: An RGB uint8 image from rendering the environment.
      extra_model_outputs: Optional. A dictionary mapping agent ids to their
          corresponding specific model outputs (also in a dictionary; e.g.
          `vf_preds` for PPO).
code: def add_env_step(
      self,
      observations: MultiAgentDict,
      actions: MultiAgentDict,
      rewards: MultiAgentDict,
      *,
      infos: Optional[MultiAgentDict] = None,
      terminateds: Optional[MultiAgentDict] = None,
      truncateds: Optional[MultiAgentDict] = None,
      render_image: Optional[np.ndarray] = None,
      extra_model_outputs: Optional[MultiAgentDict] = None,
  ) -> None:
      """Adds a timestep to the episode.
  
      Args:
          observations: A dictionary mapping agent ids to their corresponding
              observations. Note that some agents may not have stepped at this
              timestep.
          actions: Mandatory. A dictionary mapping agent ids to their
              corresponding actions. Note that some agents may not have stepped at
              this timestep.
          rewards: Mandatory. A dictionary mapping agent ids to their
              corresponding observations. Note that some agents may not have stepped
              at this timestep.
          infos: A dictionary mapping agent ids to their
              corresponding info. Note that some agents may not have stepped at this
              timestep.
          terminateds: A dictionary mapping agent ids to their `terminated` flags,
              indicating, whether the environment has been terminated for them.
              A special `__all__` key indicates that the episode is terminated for
              all agent ids.
          terminateds: A dictionary mapping agent ids to their `truncated` flags,
              indicating, whether the environment has been truncated for them.
              A special `__all__` key indicates that the episode is `truncated` for
              all agent ids.
          render_image: An RGB uint8 image from rendering the environment.
          extra_model_outputs: Optional. A dictionary mapping agent ids to their
              corresponding specific model outputs (also in a dictionary; e.g.
              `vf_preds` for PPO).
      """
      # Cannot add data to an already done episode.
      assert not self.is_done
  
      terminateds = terminateds or {}
      truncateds = truncateds or {}
  
      # Environment step.
      self.t += 1
  
      # TODO (sven, simon): Will there still be an `__all__` that is
      #  terminated or truncated?
      # TODO (simon): Maybe allow user to not provide this and then `__all__` is
      #  False?
      self.is_terminated = terminateds.get("__all__", False)
      self.is_truncated = truncateds.get("__all__", False)
  
      # Note that we store the render images into the `MultiAgentEpisode`
      # instead of storing them into each `SingleAgentEpisode`.
      if render_image is not None:
          self.render_images.append(render_image)
  
      # Add data to agent episodes.
      for agent_id in self._agent_ids:
          # Skip agents that have been terminated or truncated.
          if self.agent_episodes[agent_id].is_done:
              continue
  
          agent_is_terminated = (
              terminateds.get(agent_id, False) or self.is_terminated
          )
          agent_is_truncated = (
              truncateds.get(agent_id, False) or self.is_truncated
          )
  
          # CASE 1: observation, no action.
          # If we have an observation, but no action, we might have a buffered action,
          # or an initial agent observation.
          if agent_id in observations and agent_id not in actions:
              # We have a buffered action.
              if self.agent_buffers[agent_id]["actions"].full():
                  # Get the action from the buffer.
                  agent_action = self.agent_buffers[agent_id][
                      "actions"
                  ].get_nowait()
                  # Get the extra model output if available.
                  agent_extra_model_output = self.agent_buffers[
                      agent_id
                  ]["extra_model_outputs"].get_nowait()
                  # Reset the buffer to default value.
                  self.agent_buffers[agent_id][
                      "extra_model_outputs"
                  ].put_nowait(None)
  
                  # Note, the reward buffer will be always full with at least the
                  # default of zero reward.
                  agent_reward = self.agent_buffers[agent_id][
                      "rewards"
                  ].get_nowait()
                  # We might also got some reward in this episode.
                  if agent_id in rewards:
                      agent_reward += rewards[agent_id]
                      # Also add to the global reward list.
                      self.partial_rewards[agent_id].append(
                          rewards[agent_id]
                      )
                      # And add to the global reward timestep mapping.
                      self.partial_rewards_t[agent_id].append(self.t)
  
                  # Refill the reward buffer with the default value of zero.
                  self.agent_buffers[agent_id]["rewards"].put_nowait(
                      0.0
                  )
  
                  # It could be the last timestep of the agent.
                  # Note, in this case the agent is in `is_terminated` and
                  # `is_truncated` b/c it has an observation.
                  if agent_is_terminated or agent_is_truncated:
                      # Then, flush the buffers.
                      # TODO (simon): You can simply not refill them above.
                      self.agent_buffers[agent_id][
                          "rewards"
                      ].get_nowait()
                      self.agent_buffers[agent_id][
                          "extra_model_outputs"
                      ].get_nowait()
                  # If the agent stepped we need to keep track in the timestep
                  # mapping.
                  self.global_t_to_local_t[agent_id].append(self.t)
                  # Add data to `SingleAgentEpisode.
                  self.agent_episodes[agent_id].add_env_step(
                      observation=observations[agent_id],
                      action=agent_action,
                      reward=agent_reward,
                      infos=infos.get(agent_id),
                      terminated=agent_is_terminated,
                      truncated=agent_is_truncated,
                      extra_model_outputs=agent_extra_model_output,
                  )
              # We have no buffered action.
              else:
                  # The agent might have been terminated/truncated.
                  # if agent_id in is_terminated or agent_id in is_truncated:
                  if agent_is_terminated or agent_is_truncated:
                      # If the agent has never stepped, we treat it as not being
                      # part of this episode.
                      # Delete all of the agent's registers.
                      # del self._agent_ids[self._agent_ids.index(agent_id)]
                      # self._agent_ids.discard(agent_id)
                      # del self.agent_episodes[agent_id]
                      # del self.agent_buffers[agent_id]
                      # del self.global_t_to_local_t[agent_id]
                      # del self.global_actions_t[agent_id]
                      # del self.partial_rewards[agent_id]
                      # del self.partial_rewards_t[agent_id]
                      # Then continue with the next agent.
                      continue
                  # Then this must be the agent's initial observation.
                  else:
                      # If this was the agent's first step, record the step in the
                      # global timestep mapping.
                      self.global_t_to_local_t[agent_id].append(self.t)
                      # The agent might have got a reward.
                      # TODO (simon): Refactor to a function `record_rewards`.
                      if agent_id in rewards:
                          # Add the reward to the one in the buffer.
                          self.agent_buffers[agent_id][
                              "rewards"
                          ].put_nowait(
                              self.agent_buffers[agent_id][
                                  "rewards"
                              ].get_nowait()
                              + rewards[agent_id]
                          )
                          # Add the reward to the partial rewards of this agent.
                          self.partial_rewards[agent_id].append(
                              rewards[agent_id]
                          )
                          self.partial_rewards_t[agent_id].append(
                              self.t
                          )
  
                      self.agent_episodes[agent_id].add_env_reset(
                          observation=observations[agent_id],
                          infos=infos.get(agent_id),
                      )
          # CASE 2: No observation, but action.
          # We have no observation, but we have an action. This must be an orphane
          # action and we need to buffer it.
          elif agent_id not in observations and agent_id in actions:
              # Maybe the agent got terminated.
              if agent_is_terminated or agent_is_truncated:
                  # If this was indeed the agent's last step, we need to record it
                  # in the timestep mapping.
                  self.global_t_to_local_t[agent_id].append(self.t)
                  # Also flush all default values from the buffers.
                  self.agent_buffers[agent_id]["rewards"].get_nowait()
                  self.agent_buffers[agent_id][
                      "extra_model_outputs"
                  ].get_nowait()
                  # TODO (simon): When agents die, shall we remove them from the
                  # agent_id list?
                  # If the agent was terminated and no observation is provided,
                  # take the last one.
                  self.agent_episodes[agent_id].add_env_step(
                      observation=self.agent_episodes[
                          agent_id
                      ].observations[-1],
                      action=actions[agent_id],
                      reward=0.0
                      if agent_id not in rewards
                      else rewards[agent_id],
                      infos=self.agent_episodes[agent_id].infos[-1],
                      terminated=agent_is_terminated,
                      truncated=agent_is_truncated,
                      extra_model_outputs=None
                      if agent_id not in extra_model_outputs
                      else extra_model_outputs[agent_id],
                  )
              # Agent is still alive.
              else:
                  # TODO (simon): Maybe add a shift mapping that keeps track on
                  # original action timestep (global one). Right now the
                  # `global_reward_t` might serve here.
                  # Buffer the action.
                  self.agent_buffers[agent_id]["actions"].put_nowait(
                      actions[agent_id]
                  )
                  # Record the timestep for the action.
                  self.global_actions_t[agent_id].append(self.t)
                  # If available, buffer also reward. Note, if the agent is terminated
                  # or truncated, we finish the `SingleAgentEpisode`.
                  if agent_id in rewards:
                      # Add the reward to the existing one in the buffer. Note, the
                      # default value is zero.
                      # TODO (simon): Refactor to `record_rewards()`.
                      self.agent_buffers[agent_id][
                          "rewards"
                      ].put_nowait(
                          self.agent_buffers[agent_id][
                              "rewards"
                          ].get_nowait()
                          + rewards[agent_id]
                      )
                      # Add to the global reward list.
                      self.partial_rewards[agent_id].append(
                          rewards[agent_id]
                      )
                      # Add also to the global reward timestep mapping.
                      self.partial_rewards_t[agent_id].append(self.t)
                  # If the agent got any extra model outputs, buffer them, too.
                  if (
                      extra_model_outputs
                      and agent_id in extra_model_outputs
                  ):
                      # Flush the default `None` from buffer.
                      self.agent_buffers[agent_id][
                          "extra_model_outputs"
                      ].get_nowait()
                      # STore the extra model outputs into the buffer.
                      self.agent_buffers[agent_id][
                          "extra_model_outputs"
                      ].put_nowait(extra_model_outputs[agent_id])
  
          # CASE 3: No observation and no action.
          # We have neither observation nor action. Then, we could have `reward`,
          # `is_terminated` or `is_truncated` and should record it.
          elif agent_id not in observations and agent_id not in actions:
              # The agent could be is_terminated
              if agent_is_terminated or agent_is_truncated:
                  # If the agent has never stepped, we treat it as not being
                  # part of this episode.
                  if (
                      len(self.agent_episodes[agent_id].observations)
                      == 0
                  ):
                      # Delete all of the agent's registers.
                      # del self._agent_ids[self._agent_ids.index(agent_id)]
                      # self._agent_ids.discard(agent_id)
                      # del self.agent_episodes[agent_id]
                      # del self.agent_buffers[agent_id]
                      # del self.global_t_to_local_t[agent_id]
                      # del self.global_actions_t[agent_id]
                      # del self.partial_rewards[agent_id]
                      # del self.partial_rewards_t[agent_id]
                      # Then continue with the next agent.
                      continue
  
                  # If no observation and no action is available and the agent had
                  # stepped before the buffer must be full b/c after each
                  # observation the agent does step and if no observation followed
                  # to write the record away, the action gets buffered.
                  agent_action = self.agent_buffers[agent_id][
                      "actions"
                  ].get_nowait()
                  # Note, this is initialized as `None` if agents do not have
                  # extra model outputs.
                  agent_extra_model_output = self.agent_buffers[
                      agent_id
                  ]["extra_model_outputs"].get_nowait()
                  # Get the reward from the buffer. Note, this is always available
                  # as it is initialized as a zero reward.
                  agent_reward = self.agent_buffers[agent_id][
                      "rewards"
                  ].get_nowait()
                  # If a reward is received at this timestep record it.
                  if agent_id in rewards:
                      # TODO (simon): Refactor to `record_rewards()`.
                      agent_reward += rewards[agent_id]
                      # Add to the global reward list.
                      self.partial_rewards[agent_id].append(
                          rewards[agent_id]
                      )
                      # Add also to the global reward timestep mapping.
                      self.partial_rewards_t[agent_id].append(self.t)
  
                  # If this was indeed the agent's last step, we need to record
                  # it in the timestep mapping.
                  self.global_t_to_local_t[agent_id].append(self.t)
                  # Finish the agent's episode.
                  self.agent_episodes[agent_id].add_env_step(
                      observation=self.agent_episodes[
                          agent_id
                      ].observations[-1],
                      action=agent_action,
                      reward=agent_reward,
                      infos=infos.get(agent_id),
                      terminated=agent_is_terminated,
                      truncated=agent_is_truncated,
                      extra_model_outputs=agent_extra_model_output,
                  )
              # The agent is still alive.
              else:
                  # If the agent received an reward (triggered by actions of
                  # other agents) we collect it and add it to the one in the
                  # buffer.
                  if agent_id in rewards:
                      self.agent_buffers[agent_id][
                          "rewards"
                      ].put_nowait(
                          self.agent_buffers[agent_id][
                              "rewards"
                          ].get_nowait()
                          + rewards[agent_id]
                      )
                      # Add to the global reward list.
                      self.partial_rewards[agent_id].append(
                          rewards[agent_id]
                      )
                      # Add also to the global reward timestep mapping.
                      self.partial_rewards_t[agent_id].append(self.t)
          # CASE 4: Observation and action.
          # We have an observation and an action. Then we can simply add the
          # complete information to the episode.
          else:
              # In this case the agent id is also in the `is_terminated` and
              # `is_truncated` dictionaries.
              if agent_is_terminated or agent_is_truncated:
                  # If the agent is also done in this timestep, flush the default
                  # values from the buffers. Note, the agent has an action, i.e.
                  # it has stepped before, so we have no recorded partial rewards
                  # to add here.
                  self.agent_buffers[agent_id]["rewards"].get_nowait()
                  self.agent_buffers[agent_id][
                      "extra_model_outputs"
                  ].get_nowait()
              # If the agent stepped we need to keep track in the timestep mapping.
              self.global_t_to_local_t[agent_id].append(self.t)
              # Record the action to the global action timestep mapping.
              self.global_actions_t[agent_id].append(self.t)
              if agent_id in rewards:
                  # Also add to the global reward list.
                  self.partial_rewards[agent_id].append(
                      rewards[agent_id]
                  )
                  # And add to the global reward timestep mapping.
                  self.partial_rewards_t[agent_id].append(self.t)
              # Add timestep to `SingleAgentEpisode`.
              self.agent_episodes[agent_id].add_env_step(
                  observation=observations[agent_id],
                  action=actions[agent_id],
                  reward=0.0
                  if agent_id not in rewards
                  else rewards[agent_id],
                  infos=infos.get(agent_id),
                  terminated=agent_is_terminated,
                  truncated=agent_is_truncated,
                  extra_model_outputs=None
                  if extra_model_outputs is None
                  else extra_model_outputs[agent_id],
              )
  
imports: from collections import defaultdict
  from queue import Queue
  from typing import Any, Dict, List, Optional, Set, Union
  import uuid
  import numpy
  from ray.rllib.env.single_agent_episode import SingleAgentEpisode
  from ray.rllib.policy.sample_batch import MultiAgentBatch
  from ray.rllib.utils.typing import MultiAgentDict
label_desc: Use image node_id 12 for calling the MultiAgentEpisode obj's underlying member method code with example usage: obj.add_env_step(observations, actions, rewards) without return types