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))
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.
    
tensor2tensor/rl/dopamine_connector.py