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.
  """

  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.
    """
    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)

  return create_agent