research/efficient-hrl/train_utils.py
# 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)