class_docstr: Tracks the current state of a (possibly multi-agent) episode.
  
  Attributes:
      new_batch_builder: Create a new MultiAgentSampleBatchBuilder.
      add_extra_batch: Return a built MultiAgentBatch to the sampler.
      batch_builder: Batch builder for the current episode.
      total_reward: Summed reward across all agents in this episode.
      length: Length of this episode.
      episode_id: Unique id identifying this trajectory.
      agent_rewards: Summed rewards broken down by agent.
      custom_metrics: Dict where the you can add custom metrics.
      user_data: Dict that you can use for temporary storage. E.g.
          in between two custom callbacks referring to the same episode.
      hist_data: Dict mapping str keys to List[float] for storage of
          per-timestep float data throughout the episode.
  
  Use case 1: Model-based rollouts in multi-agent:
      A custom compute_actions() function in a policy can inspect the
      current episode state and perform a number of rollouts based on the
      policies and state of other agents in the environment.
  
  Use case 2: Returning extra rollouts data.
      The model rollouts can be returned back to the sampler by calling:
  
      .. testcode::
          :skipif: True
  
          batch = episode.new_batch_builder()
          for each transition:
              batch.add_values(...)  # see sampler for usage
          episode.extra_batches.add(batch.build_and_reset())
docstr: Initializes an Episode instance.
  
  Args:
      policies: The PolicyMap object (mapping PolicyIDs to Policy
          objects) to use for determining, which policy is used for
          which agent.
      policy_mapping_fn: The mapping function mapping AgentIDs to
          PolicyIDs.
      batch_builder_factory:
      extra_batch_callback:
      env_id: The environment's ID in which this episode runs.
      worker: The RolloutWorker instance, in which this episode runs.
docstr_tok: ["Initializes","an","Episode","instance",".","Args",":","policies",":","The","PolicyMap","object","(","mapping","PolicyIDs","to","Policy","objects",")","to","use","for","determining",",","which","policy","is","used","for","which","agent",".","policy_mapping_fn",":","The","mapping","function","mapping","AgentIDs","to","PolicyIDs",".","batch_builder_factory",":","extra_batch_callback",":","env_id",":","The","environment","'s","ID","in","which","this","episode","runs",".","worker",":","The","RolloutWorker","instance",",","in","which","this","episode","runs","."]
code: def __init__(
      self,
      policies: PolicyMap,
      policy_mapping_fn: Callable[
          [AgentID, "Episode", "RolloutWorker"], PolicyID
      ],
      batch_builder_factory: Callable[
          [], "MultiAgentSampleBatchBuilder"
      ],
      extra_batch_callback: Callable[[SampleBatchType], None],
      env_id: EnvID,
      *,
      worker: Optional["RolloutWorker"] = None,
  ):
      """Initializes an Episode instance.
  
      Args:
          policies: The PolicyMap object (mapping PolicyIDs to Policy
              objects) to use for determining, which policy is used for
              which agent.
          policy_mapping_fn: The mapping function mapping AgentIDs to
              PolicyIDs.
          batch_builder_factory:
          extra_batch_callback:
          env_id: The environment's ID in which this episode runs.
          worker: The RolloutWorker instance, in which this episode runs.
      """
      self.new_batch_builder: Callable[
          [], "MultiAgentSampleBatchBuilder"
      ] = batch_builder_factory
      self.add_extra_batch: Callable[
          [SampleBatchType], None
      ] = extra_batch_callback
      self.batch_builder: "MultiAgentSampleBatchBuilder" = (
          batch_builder_factory()
      )
      self.total_reward: float = 0.0
      self.length: int = 0
      self.started = False
      self.episode_id: int = random.randrange(int(1e18))
      self.env_id = env_id
      self.worker = worker
      self.agent_rewards: Dict[
          Tuple[AgentID, PolicyID], float
      ] = defaultdict(float)
      self.custom_metrics: Dict[str, float] = {}
      self.user_data: Dict[str, Any] = {}
      self.hist_data: Dict[str, List[float]] = {}
      self.media: Dict[str, Any] = {}
      self.policy_map: PolicyMap = policies
      self._policies = self.policy_map  # backward compatibility
      self.policy_mapping_fn: Callable[
          [AgentID, "Episode", "RolloutWorker"], PolicyID
      ] = policy_mapping_fn
      self.is_faulty = False
      self._next_agent_index: int = 0
      self._agent_to_index: Dict[AgentID, int] = {}
      self._agent_to_policy: Dict[AgentID, PolicyID] = {}
      self._agent_to_rnn_state: Dict[AgentID, List[Any]] = {}
      self._agent_to_last_obs: Dict[AgentID, EnvObsType] = {}
      self._agent_to_last_raw_obs: Dict[AgentID, EnvObsType] = {}
      self._agent_to_last_terminated: Dict[AgentID, bool] = {}
      self._agent_to_last_truncated: Dict[AgentID, bool] = {}
      self._agent_to_last_info: Dict[AgentID, EnvInfoDict] = {}
      self._agent_to_last_action: Dict[AgentID, EnvActionType] = {}
      self._agent_to_last_extra_action_outs: Dict[AgentID, dict] = {}
      self._agent_to_prev_action: Dict[AgentID, EnvActionType] = {}
      self._agent_reward_history: Dict[
          AgentID, List[int]
      ] = defaultdict(list)
  
imports: import random
  from collections import defaultdict
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
  import numpy
  import tree
  from ray.rllib.env.base_env import _DUMMY_AGENT_ID
  from ray.rllib.policy.policy_map import PolicyMap
  from ray.rllib.utils.annotations import Deprecated, DeveloperAPI
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray
  from ray.rllib.utils.typing import AgentID, EnvActionType, EnvID, EnvInfoDict, EnvObsType, PolicyID, SampleBatchType
  from ray.util import log_once
label_desc: Use image node_id 1 to create a new Episode object  with example: obj = Episode(policies, policy_mapping_fn, batch_builder_factory, extra_batch_callback, env_id)