google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
r"""Train an agent.
17

18
"""
19
import os
20

21

22
from absl import app
23
from absl import flags
24

25
from dopamine.discrete_domains import run_experiment
26
import tensorflow.compat.v1 as tf
27

28
from experience_replay import run_experience_replay_experiment
29

30

31
flags.DEFINE_string('base_dir', None,
32
                    'Base directory to host all required sub-directories.')
33
flags.DEFINE_multi_string(
34
    'gin_files', [], 'List of paths to gin configuration files (e.g.'
35
    '"third_party/py/dopamine/agents/dqn/dqn.gin").')
36
flags.DEFINE_multi_string(
37
    'gin_bindings', [],
38
    'Gin bindings to override the values set in the config files '
39
    '(e.g. "DQNAgent.epsilon_train=0.1",'
40
    '      "create_atari_environment.game_name="Pong"").')
41
flags.DEFINE_string(
42
    'schedule', 'continuous_train_and_eval',
43
    'The schedule with which to run the experiment and choose an appropriate '
44
    'Runner. Supported choices are '
45
    '{continuous_train, eval, continuous_train_and_eval}.')
46

47

48
FLAGS = flags.FLAGS
49

50

51

52
def create_runner(base_dir, create_agent_fn,
53
                  schedule='continuous_train_and_eval'):
54
  """Creates an experiment Runner.
55

56
  TODO(b/): Figure out the right idiom to create a Runner. The current mechanism
57
  of using a number of flags will not scale and is not elegant.
58

59
  Args:
60
    base_dir: Base directory for hosting all subdirectories.
61
    create_agent_fn: A function that takes as args a Tensorflow session and a
62
     Gym Atari 2600 environment, and returns an agent.
63
    schedule: string, which type of Runner to use.
64

65
  Returns:
66
    runner: A `run_experiment.Runner` like object.
67

68
  Raises:
69
    ValueError: When an unknown schedule is encountered.
70
  """
71
  assert base_dir is not None
72

73
  # Continuously runs training and eval till max num_iterations is hit.
74
  if schedule == 'continuous_train_and_eval':
75
    return run_experience_replay_experiment.ElephantRunner(
76
        base_dir, create_agent_fn)
77

78
  else:
79
    raise ValueError('Unknown schedule: {}'.format(schedule))
80

81

82
def launch_experiment(create_runner_fn, create_agent_fn):
83
  """Launches the experiment.
84

85
  Args:
86
    create_runner_fn: A function that takes as args a base directory and a
87
      function for creating an agent and returns a `Runner` like object.
88
    create_agent_fn: A function that takes as args a Tensorflow session and a
89
      Gym environment, and returns an agent.
90
  """
91
  run_experiment.load_gin_configs(FLAGS.gin_files, FLAGS.gin_bindings)
92
  runner = create_runner_fn(FLAGS.base_dir, create_agent_fn,
93
                            schedule=FLAGS.schedule)
94
  runner.run_experiment()
95

96

97
def main(unused_argv):
98
  """This main function acts as a wrapper around a gin-configurable experiment.
99

100
  Args:
101
    unused_argv: Arguments (unused).
102
  """
103
  tf.logging.set_verbosity(tf.logging.INFO)
104
  launch_experiment(create_runner,
105
                    run_experience_replay_experiment.create_agent)
106

107
if __name__ == '__main__':
108
  flags.mark_flag_as_required('base_dir')
109
  app.run(main)
110