class_docstr: A JAX implementation of the Offline DQN agent with DR3 regularizer.
class_docstr_tok: ["A","JAX","implementation","of","the","Offline","DQN","agent","with","DR3","regularizer","."]
docstr: Initializes the agent and constructs the necessary components.
  
  Args:
    num_actions: int, number of actions the agent can take at any state.
    replay_data_dir: str, log Directory from which to load the replay buffer.
    dr3_coefficient: float, Coefficient for the DR3 regularizer.
    cql_coefficient: float, Coefficient for the CQL loss.
    summary_writer: SummaryWriter object for outputting training statistics.
    replay_buffer_builder: Callable object that takes "self" as an argument
      and returns a replay buffer to use for training offline. If None,
      it will use the default FixedReplayBuffer.
docstr_tok: ["Initializes","the","agent","and","constructs","the","necessary","components",".","Args",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","replay_data_dir",":","str",",","log","Directory","from","which","to","load","the","replay","buffer",".","dr3_coefficient",":","float",",","Coefficient","for","the","DR3","regularizer",".","cql_coefficient",":","float",",","Coefficient","for","the","CQL","loss",".","summary_writer",":","SummaryWriter","object","for","outputting","training","statistics",".","replay_buffer_builder",":","Callable","object","that","takes","``","self","''","as","an","argument","and","returns","a","replay","buffer","to","use","for","training","offline",".","If","None",",","it","will","use","the","default","FixedReplayBuffer","."]
code: def __init__(
      self,
      num_actions,
      replay_data_dir,
      dr3_coefficient=0.0,
      cql_coefficient=0.0,
      summary_writer=None,
      replay_buffer_builder=None,
  ):
      """Initializes the agent and constructs the necessary components.
  
      Args:
        num_actions: int, number of actions the agent can take at any state.
        replay_data_dir: str, log Directory from which to load the replay buffer.
        dr3_coefficient: float, Coefficient for the DR3 regularizer.
        cql_coefficient: float, Coefficient for the CQL loss.
        summary_writer: SummaryWriter object for outputting training statistics.
        replay_buffer_builder: Callable object that takes "self" as an argument
          and returns a replay buffer to use for training offline. If None,
          it will use the default FixedReplayBuffer.
      """
      logging.info(
          "Creating %s agent with the following parameters:",
          self.__class__.__name__,
      )
      logging.info("\t dr3_coefficient: %s", dr3_coefficient)
      logging.info("\t CQL coefficient: %s", cql_coefficient)
      super().__init__(
          num_actions,
          replay_data_dir=replay_data_dir,
          replay_buffer_builder=replay_buffer_builder,
          summary_writer=summary_writer,
      )
      self._dr3_coefficient = dr3_coefficient
      self._cql_coefficient = cql_coefficient
  
file_path: dopamine/dopamine/labs/offline_rl/jax/offline_dr3_agent.py
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.labs.offline_rl.jax import offline_dqn_agent
  from dopamine.metrics import statistics_instance
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 1 to create a new OfflineJaxDR3Agent object from inherited base classes: offline_dqn_agent with example: obj = OfflineJaxDR3Agent(num_actions, replay_data_dir, dr3_coefficient, cql_coefficient, summary_writer, replay_buffer_builder)