File size: 7,027 Bytes
0b8359d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

r""""""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import namedtuple
import os
import time

import tensorflow as tf

import gin.tf

flags = tf.app.flags


flags.DEFINE_multi_string('config_file', None,
                          'List of paths to the config files.')
flags.DEFINE_multi_string('params', None,
                          'Newline separated list of Gin parameter bindings.')

flags.DEFINE_string('train_dir', None,
                    'Directory for writing logs/summaries during training.')
flags.DEFINE_string('master', 'local',
                    'BNS name of the TensorFlow master to use.')
flags.DEFINE_integer('task', 0, 'task id')
flags.DEFINE_integer('save_interval_secs', 300, 'The frequency at which '
                     'checkpoints are saved, in seconds.')
flags.DEFINE_integer('save_summaries_secs', 30, 'The frequency at which '
                     'summaries are saved, in seconds.')
flags.DEFINE_boolean('summarize_gradients', False,
                     'Whether to generate gradient summaries.')

FLAGS = flags.FLAGS

TrainOps = namedtuple('TrainOps',
                      ['train_op', 'meta_train_op', 'collect_experience_op'])


class TrainStep(object):
  """Handles training step."""

  def __init__(self,
               max_number_of_steps=0,
               num_updates_per_observation=1,
               num_collect_per_update=1,
               num_collect_per_meta_update=1,
               log_every_n_steps=1,
               policy_save_fn=None,
               save_policy_every_n_steps=0,
               should_stop_early=None):
    """Returns a function that is executed at each step of slim training.

    Args:
      max_number_of_steps: Optional maximum number of train steps to take.
      num_updates_per_observation: Number of updates per observation.
      log_every_n_steps: The frequency, in terms of global steps, that the loss
      and global step and logged.
      policy_save_fn: A tf.Saver().save function to save the policy.
      save_policy_every_n_steps: How frequently to save the policy.
      should_stop_early: Optional hook to report whether training should stop.
    Raises:
      ValueError: If policy_save_fn is not provided when
        save_policy_every_n_steps > 0.
    """
    if save_policy_every_n_steps and policy_save_fn is None:
      raise ValueError(
          'policy_save_fn is required when save_policy_every_n_steps > 0')
    self.max_number_of_steps = max_number_of_steps
    self.num_updates_per_observation = num_updates_per_observation
    self.num_collect_per_update = num_collect_per_update
    self.num_collect_per_meta_update = num_collect_per_meta_update
    self.log_every_n_steps = log_every_n_steps
    self.policy_save_fn = policy_save_fn
    self.save_policy_every_n_steps = save_policy_every_n_steps
    self.should_stop_early = should_stop_early
    self.last_global_step_val = 0
    self.train_op_fn = None
    self.collect_and_train_fn = None
    tf.logging.info('Training for %d max_number_of_steps',
                    self.max_number_of_steps)

  def train_step(self, sess, train_ops, global_step, _):
    """This function will be called at each step of training.

    This represents one step of the DDPG algorithm and can include:
    1. collect a <state, action, reward, next_state> transition
    2. update the target network
    3. train the actor
    4. train the critic

    Args:
      sess: A Tensorflow session.
      train_ops: A DdpgTrainOps tuple of train ops to run.
      global_step: The global step.

    Returns:
      A scalar total loss.
      A boolean should stop.
    """
    start_time = time.time()
    if self.train_op_fn is None:
      self.train_op_fn = sess.make_callable([train_ops.train_op, global_step])
      self.meta_train_op_fn = sess.make_callable([train_ops.meta_train_op, global_step])
      self.collect_fn = sess.make_callable([train_ops.collect_experience_op, global_step])
      self.collect_and_train_fn = sess.make_callable(
          [train_ops.train_op, global_step, train_ops.collect_experience_op])
      self.collect_and_meta_train_fn = sess.make_callable(
          [train_ops.meta_train_op, global_step, train_ops.collect_experience_op])
    for _ in range(self.num_collect_per_update - 1):
      self.collect_fn()
    for _ in range(self.num_updates_per_observation - 1):
      self.train_op_fn()

    total_loss, global_step_val, _ = self.collect_and_train_fn()
    if (global_step_val // self.num_collect_per_meta_update !=
        self.last_global_step_val // self.num_collect_per_meta_update):
      self.meta_train_op_fn()

    time_elapsed = time.time() - start_time
    should_stop = False
    if self.max_number_of_steps:
      should_stop = global_step_val >= self.max_number_of_steps
    if global_step_val != self.last_global_step_val:
      if (self.save_policy_every_n_steps and
          global_step_val // self.save_policy_every_n_steps !=
          self.last_global_step_val // self.save_policy_every_n_steps):
        self.policy_save_fn(sess)

      if (self.log_every_n_steps and
          global_step_val % self.log_every_n_steps == 0):
        tf.logging.info(
            'global step %d: loss = %.4f (%.3f sec/step) (%d steps/sec)',
            global_step_val, total_loss, time_elapsed, 1 / time_elapsed)

    self.last_global_step_val = global_step_val
    stop_early = bool(self.should_stop_early and self.should_stop_early())
    return total_loss, should_stop or stop_early


def create_counter_summaries(counters):
  """Add named summaries to counters, a list of tuples (name, counter)."""
  if counters:
    with tf.name_scope('Counters/'):
      for name, counter in counters:
        tf.summary.scalar(name, counter)


def gen_debug_batch_summaries(batch):
  """Generates summaries for the sampled replay batch."""
  states, actions, rewards, _, next_states = batch
  with tf.name_scope('batch'):
    for s in range(states.get_shape()[-1]):
      tf.summary.histogram('states_%d' % s, states[:, s])
    for s in range(states.get_shape()[-1]):
      tf.summary.histogram('next_states_%d' % s, next_states[:, s])
    for a in range(actions.get_shape()[-1]):
      tf.summary.histogram('actions_%d' % a, actions[:, a])
    tf.summary.histogram('rewards', rewards)