class_docstr: Replay Buffers for loading existing data.
class_docstr_tok: ["Replay","Buffers","for","loading","existing","data","."]
docstr: Initialize the JaxMultiTaskFixedReplayBuffer class.
  
  Args:
    replay_capacity: int, number of transitions to keep in memory. This can be
      used with `replay_transitions_start_index` to read a subset of replay
      data starting from a specific position.
    batch_size: int, Batch size for sampling data from buffer.
    dataset_name: str, Name of the game to load data from.
    dataset_expertise: Optional[float], expertise of the dataset to use.
    stack_size: int, number of frames to use in state stack.
    update_horizon: int, length of update ('n' in n-step update).
    gamma: int, the discount factor.
    return_to_go: bool, whether to add return_to_go to the replay dataset.
    unused_kwargs: Unused kwargs arguments.
docstr_tok: ["Initialize","the","JaxMultiTaskFixedReplayBuffer","class",".","Args",":","replay_capacity",":","int",",","number","of","transitions","to","keep","in","memory",".","This","can","be","used","with","`","replay_transitions_start_index","`","to","read","a","subset","of","replay","data","starting","from","a","specific","position",".","batch_size",":","int",",","Batch","size","for","sampling","data","from","buffer",".","dataset_name",":","str",",","Name","of","the","game","to","load","data","from",".","dataset_expertise",":","Optional","[","float","]",",","expertise","of","the","dataset","to","use",".","stack_size",":","int",",","number","of","frames","to","use","in","state","stack",".","update_horizon",":","int",",","length","of","update","(","'n","'","in","n-step","update",")",".","gamma",":","int",",","the","discount","factor",".","return_to_go",":","bool",",","whether","to","add","return_to_go","to","the","replay","dataset",".","unused_kwargs",":","Unused","kwargs","arguments","."]
code: def __init__(
      self,
      replay_capacity,
      batch_size,
      dataset_name,
      dataset_expertise=None,
      stack_size=4,
      update_horizon=1,
      gamma=0.99,
      return_to_go=False,
      **unused_kwargs,
  ):
      """Initialize the JaxMultiTaskFixedReplayBuffer class.
  
      Args:
        replay_capacity: int, number of transitions to keep in memory. This can be
          used with `replay_transitions_start_index` to read a subset of replay
          data starting from a specific position.
        batch_size: int, Batch size for sampling data from buffer.
        dataset_name: str, Name of the game to load data from.
        dataset_expertise: Optional[float], expertise of the dataset to use.
        stack_size: int, number of frames to use in state stack.
        update_horizon: int, length of update ('n' in n-step update).
        gamma: int, the discount factor.
        return_to_go: bool, whether to add return_to_go to the replay dataset.
        unused_kwargs: Unused kwargs arguments.
      """
  
      logging.info(
          "Creating JaxFixedReplayBufferTFDS with the parameters:"
      )
      logging.info("\t Batch size: %d", batch_size)
      logging.info("\t Update horizon: %d", update_horizon)
      logging.info("\t Gamma: %f", gamma)
      logging.info("\t Return To Go: %s", gamma)
      logging.info("\t TFDS dataset: %s", dataset_name)
      logging.info("\t Dataset expertise: %s", dataset_expertise)
      # data % = (replay_capacity / 1M) x 100
      self._data_percent = float(replay_capacity) / 10_000.0
      self._replay_capacity = replay_capacity
      self._batch_size = batch_size
      self._return_to_go = return_to_go
      self._total_steps = self._data_percent * 50_000_000
      self._dataset_expertise = dataset_expertise
      if dataset_expertise is not None:
          dataset, (
              self._total_steps,
              self._min_score,
              self._max_score,
          ) = scaling_dataset_utils.create_dataset_with_expertise(
              game=game_from_dataset_name(dataset_name),
              dataset_expertise=dataset_expertise,
          )
      else:
          dataset = tfds_atari_utils.uniformly_subsampled_atari_data(
              dataset_name, self._data_percent
          )
      logging.info("\t StepsPerEpoch %d", self.gradient_steps_per_epoch)
      self._replay = tfds_atari_utils.build_tfds_replay(
          dataset=dataset,
          stack_size=stack_size,
          update_horizon=update_horizon,
          gamma=gamma,
          batch_size=batch_size,
          return_to_go=return_to_go,
      )
  
imports: import os
  from absl import logging
  from dopamine.labs.offline_rl.rlu_tfds import scaling_dataset_utils
  from dopamine.labs.offline_rl.rlu_tfds import tfds_atari_utils
  import gin
label_desc: Use image node_id 1 to create a new JaxFixedReplayBufferTFDS object from inherited base classes: object with example: obj = JaxFixedReplayBufferTFDS(replay_capacity, batch_size, dataset_name, dataset_expertise, stack_size, update_horizon, gamma, return_to_go)