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())
class_docstr_tok: ["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","(",")",")"]
code: def _set_last_action(self, agent_id, action):
      if agent_id in self._agent_to_last_action:
          self._agent_to_prev_action[
              agent_id
          ] = self._agent_to_last_action[agent_id]
      self._agent_to_last_action[agent_id] = action
  
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 22 for calling the Episode obj's underlying member method code with example usage: obj._set_last_action(agent_id, action) without return types