def get_create_agent(agent_kwargs):
  """Factory for dopamine agent initialization.

  Args:
    agent_kwargs: dict of BatchDQNAgent parameters

  Returns:
    Function(sess, environment, summary_writer) -> BatchDQNAgent instance.
  """
  agent_kwargs = copy.deepcopy(agent_kwargs)
  agent_type = agent_kwargs.pop("type")

  def create_agent(sess, environment, summary_writer=None):
    """Creates a DQN agent.

    Simplified version of `dopamine.discrete_domains.train.create_agent`

    Args:
      sess: a session
      environment: an environment
      summary_writer: a summary writer.

    Returns:
      a DQN agent.
    """
    if agent_type == "Rainbow":
      return BatchRainbowAgent(
          env_batch_size=environment.batch_size,
          sess=sess,
          num_actions=environment.action_space.n,
          summary_writer=summary_writer,
          tf_device="/gpu:*",
          **agent_kwargs)
    elif agent_type == "VanillaDQN":
      return BatchDQNAgent(
          env_batch_size=environment.batch_size,
          sess=sess,
          num_actions=environment.action_space.n,
          summary_writer=summary_writer,
          tf_device="/gpu:*",
          **agent_kwargs)
    else:
      raise ValueError("Unknown agent_type {}".format(agent_type))

  return create_agent
Factory for dopamine agent initialization.

  Args:
    agent_kwargs: dict of BatchDQNAgent parameters

  Returns:
    Function(sess, environment, summary_writer) -> BatchDQNAgent instance.
  
tensor2tensor/rl/dopamine_connector.py