docstr: :param max_size: the maximum number of transitions or episodes to hold in the memory
  :param allow_duplicates_in_batch_sampling: allow having the same transition multiple times in a batch
  :param num_classes: the number of classes in the replayed data
  :param state_key_with_the_class_index: the class index is assumed to be a value in the state dictionary.
                                     this parameter determines the key to retrieve the class index value
docstr_tok: [":","param","max_size",":","the","maximum","number","of","transitions","or","episodes","to","hold","in","the","memory",":","param","allow_duplicates_in_batch_sampling",":","allow","having","the","same","transition","multiple","times","in","a","batch",":","param","num_classes",":","the","number","of","classes","in","the","replayed","data",":","param","state_key_with_the_class_index",":","the","class","index","is","assumed","to","be","a","value","in","the","state","dictionary",".","this","parameter","determines","the","key","to","retrieve","the","class","index","value"]
code: def __init__(
      self,
      max_size: Tuple[MemoryGranularity, int],
      allow_duplicates_in_batch_sampling: bool = True,
      num_classes: int = 0,
      state_key_with_the_class_index: Any = "class",
  ):
      """
      :param max_size: the maximum number of transitions or episodes to hold in the memory
      :param allow_duplicates_in_batch_sampling: allow having the same transition multiple times in a batch
      :param num_classes: the number of classes in the replayed data
      :param state_key_with_the_class_index: the class index is assumed to be a value in the state dictionary.
                                         this parameter determines the key to retrieve the class index value
      """
      super().__init__(max_size, allow_duplicates_in_batch_sampling)
      self.current_class_to_sample_from = 0
      self.num_classes = num_classes
      self.state_key_with_the_class_index = (
          state_key_with_the_class_index
      )
      self.transitions = [[] for _ in range(self.num_classes)]
      self.transitions_order = []
  
      if self.num_classes < 2:
          raise ValueError(
              "The number of classes for a balanced replay buffer should be at least 2. "
              "The number of classes that were defined are: {}".format(
                  self.num_classes
              )
          )
  
imports: import operator
  import random
  from enum import Enum
  from typing import List, Tuple, Any, Union
  import numpy
  from rl_coach.core_types import Transition
  from rl_coach.memories.memory import MemoryGranularity
  from rl_coach.memories.non_episodic.experience_replay import ExperienceReplayParameters, ExperienceReplay
  from rl_coach.schedules import Schedule, ConstantSchedule
label_desc: Use image node_id 1 to create a new BalancedExperienceReplay object from inherited base classes: ExperienceReplay with example: obj = BalancedExperienceReplay(max_size, allow_duplicates_in_batch_sampling, num_classes, state_key_with_the_class_index)