Spaces:
Running
Running
| # 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) | |