code: def create_offline_agent(
      sess,
      environment,
      agent_name,
      replay_data_dir,
      summary_writer=None,
  ):
      """Creates an online agent.
  
      Args:
        sess: A `tf.Session` object for running associated ops. This argument is
          ignored for JAX agents.
        environment: An Atari 2600 environment.
        agent_name: Name of the agent to be created.
        replay_data_dir: Directory from which to load the fixed replay buffers.
        summary_writer: A Tensorflow summary writer to pass to the agent for
          in-agent training statistics in Tensorboard.
  
      Returns:
        An agent with metrics.
      """
      if agent_name == "jax_dqn":
          agent = offline_dqn_agent.OfflineJaxDQNAgent
      elif agent_name == "jax_dr3":
          agent = offline_dr3_agent.OfflineJaxDR3Agent
      elif agent_name == "jax_rainbow":
          agent = offline_rainbow_agent.OfflineJaxRainbowAgent
      elif agent_name == "jax_return_conditioned_bc":
          agent = (
              return_conditioned_bc_agent.JaxReturnConditionedBCAgent
          )
      elif agent_name == "jax_dr3_with_validation":
          agent = (
              offline_dr3_with_validation_agent.OfflineJaxDR3WithValidationAgent
          )
      else:
          raise ValueError(
              "{} is not a valid agent name".format(FLAGS.agent_name)
          )
  
      if agent_name.startswith("jax"):  # JAX agent
          return agent(
              num_actions=environment.action_space.n,
              replay_data_dir=replay_data_dir,
              summary_writer=summary_writer,
              use_tfds=FLAGS.use_tfds,
          )
      else:
          return agent(
              sess,
              num_actions=environment.action_space.n,
              replay_data_dir=replay_data_dir,
              summary_writer=summary_writer,
          )
  
file_path: dopamine/dopamine/labs/offline_rl/jax/train.py
imports: import functools
  import json
  import os
  from absl import app
  from absl import flags
  from absl import logging
  from dopamine.discrete_domains import run_experiment
  from dopamine.discrete_domains import train
  from dopamine.labs.offline_rl.jax import offline_dqn_agent
  from dopamine.labs.offline_rl.jax import offline_dr3_agent
  from dopamine.labs.offline_rl.jax import offline_dr3_with_validation_agent
  from dopamine.labs.offline_rl.jax import offline_rainbow_agent
  from dopamine.labs.offline_rl.jax import return_conditioned_bc_agent
  from dopamine.labs.offline_rl.jax import run_experiment
  from jax import config
label_desc: Use image node_id 1 for calling a global function with example usage: create_offline_agent(sess, environment, agent_name, replay_data_dir, summary_writer) and returns: agent, agent