# 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"""Script for training an RL agent using the UVF algorithm. To run locally: See run_train.py """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import time import tensorflow as tf slim = tf.contrib.slim import gin.tf # pylint: disable=unused-import import train_utils import agent as agent_ from agents import circular_buffer from utils import utils as uvf_utils from environments import create_maze_env # pylint: enable=unused-import flags = tf.app.flags FLAGS = flags.FLAGS flags.DEFINE_string('goal_sample_strategy', 'sample', 'None, sample, FuN') LOAD_PATH = None def collect_experience(tf_env, agent, meta_agent, state_preprocess, replay_buffer, meta_replay_buffer, action_fn, meta_action_fn, environment_steps, num_episodes, num_resets, episode_rewards, episode_meta_rewards, store_context, disable_agent_reset): """Collect experience in a tf_env into a replay_buffer using action_fn. Args: tf_env: A TFEnvironment. agent: A UVF agent. meta_agent: A Meta Agent. replay_buffer: A Replay buffer to collect experience in. meta_replay_buffer: A Replay buffer to collect meta agent experience in. action_fn: A function to produce actions given current state. meta_action_fn: A function to produce meta actions given current state. environment_steps: A variable to count the number of steps in the tf_env. num_episodes: A variable to count the number of episodes. num_resets: A variable to count the number of resets. store_context: A boolean to check if store context in replay. disable_agent_reset: A boolean that disables agent from resetting. Returns: A collect_experience_op that excute an action and store into the replay_buffers """ tf_env.start_collect() state = tf_env.current_obs() state_repr = state_preprocess(state) action = action_fn(state, context=None) with tf.control_dependencies([state]): transition_type, reward, discount = tf_env.step(action) def increment_step(): return environment_steps.assign_add(1) def increment_episode(): return num_episodes.assign_add(1) def increment_reset(): return num_resets.assign_add(1) def update_episode_rewards(context_reward, meta_reward, reset): new_episode_rewards = tf.concat( [episode_rewards[:1] + context_reward, episode_rewards[1:]], 0) new_episode_meta_rewards = tf.concat( [episode_meta_rewards[:1] + meta_reward, episode_meta_rewards[1:]], 0) return tf.group( episode_rewards.assign( tf.cond(reset, lambda: tf.concat([[0.], episode_rewards[:-1]], 0), lambda: new_episode_rewards)), episode_meta_rewards.assign( tf.cond(reset, lambda: tf.concat([[0.], episode_meta_rewards[:-1]], 0), lambda: new_episode_meta_rewards))) def no_op_int(): return tf.constant(0, dtype=tf.int64) step_cond = agent.step_cond_fn(state, action, transition_type, environment_steps, num_episodes) reset_episode_cond = agent.reset_episode_cond_fn( state, action, transition_type, environment_steps, num_episodes) reset_env_cond = agent.reset_env_cond_fn(state, action, transition_type, environment_steps, num_episodes) increment_step_op = tf.cond(step_cond, increment_step, no_op_int) increment_episode_op = tf.cond(reset_episode_cond, increment_episode, no_op_int) increment_reset_op = tf.cond(reset_env_cond, increment_reset, no_op_int) increment_op = tf.group(increment_step_op, increment_episode_op, increment_reset_op) with tf.control_dependencies([increment_op, reward, discount]): next_state = tf_env.current_obs() next_state_repr = state_preprocess(next_state) next_reset_episode_cond = tf.logical_or( agent.reset_episode_cond_fn( state, action, transition_type, environment_steps, num_episodes), tf.equal(discount, 0.0)) if store_context: context = [tf.identity(var) + tf.zeros_like(var) for var in agent.context_vars] meta_context = [tf.identity(var) + tf.zeros_like(var) for var in meta_agent.context_vars] else: context = [] meta_context = [] with tf.control_dependencies([next_state] + context + meta_context): if disable_agent_reset: collect_experience_ops = [tf.no_op()] # don't reset agent else: collect_experience_ops = agent.cond_begin_episode_op( tf.logical_not(reset_episode_cond), [state, action, reward, next_state, state_repr, next_state_repr], mode='explore', meta_action_fn=meta_action_fn) context_reward, meta_reward = collect_experience_ops collect_experience_ops = list(collect_experience_ops) collect_experience_ops.append( update_episode_rewards(tf.reduce_sum(context_reward), meta_reward, reset_episode_cond)) meta_action_every_n = agent.tf_context.meta_action_every_n with tf.control_dependencies(collect_experience_ops): transition = [state, action, reward, discount, next_state] meta_action = tf.to_float( tf.concat(context, -1)) # Meta agent action is low-level context meta_end = tf.logical_and( # End of meta-transition. tf.equal(agent.tf_context.t % meta_action_every_n, 1), agent.tf_context.t > 1) with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): states_var = tf.get_variable('states_var', [meta_action_every_n, state.shape[-1]], state.dtype) actions_var = tf.get_variable('actions_var', [meta_action_every_n, action.shape[-1]], action.dtype) state_var = tf.get_variable('state_var', state.shape, state.dtype) reward_var = tf.get_variable('reward_var', reward.shape, reward.dtype) meta_action_var = tf.get_variable('meta_action_var', meta_action.shape, meta_action.dtype) meta_context_var = [ tf.get_variable('meta_context_var%d' % idx, meta_context[idx].shape, meta_context[idx].dtype) for idx in range(len(meta_context))] actions_var_upd = tf.scatter_update( actions_var, (agent.tf_context.t - 2) % meta_action_every_n, action) with tf.control_dependencies([actions_var_upd]): actions = tf.identity(actions_var) + tf.zeros_like(actions_var) meta_reward = tf.identity(meta_reward) + tf.zeros_like(meta_reward) meta_reward = tf.reshape(meta_reward, reward.shape) reward = 0.1 * meta_reward meta_transition = [state_var, meta_action_var, reward_var + reward, discount * (1 - tf.to_float(next_reset_episode_cond)), next_state] meta_transition.extend([states_var, actions]) if store_context: # store current and next context into replay transition += context + list(agent.context_vars) meta_transition += meta_context_var + list(meta_agent.context_vars) meta_step_cond = tf.squeeze(tf.logical_and(step_cond, tf.logical_or(next_reset_episode_cond, meta_end))) collect_experience_op = tf.group( replay_buffer.maybe_add(transition, step_cond), meta_replay_buffer.maybe_add(meta_transition, meta_step_cond), ) with tf.control_dependencies([collect_experience_op]): collect_experience_op = tf.cond(reset_env_cond, tf_env.reset, tf_env.current_time_step) meta_period = tf.equal(agent.tf_context.t % meta_action_every_n, 1) states_var_upd = tf.scatter_update( states_var, (agent.tf_context.t - 1) % meta_action_every_n, next_state) state_var_upd = tf.assign( state_var, tf.cond(meta_period, lambda: next_state, lambda: state_var)) reward_var_upd = tf.assign( reward_var, tf.cond(meta_period, lambda: tf.zeros_like(reward_var), lambda: reward_var + reward)) meta_action = tf.to_float(tf.concat(agent.context_vars, -1)) meta_action_var_upd = tf.assign( meta_action_var, tf.cond(meta_period, lambda: meta_action, lambda: meta_action_var)) meta_context_var_upd = [ tf.assign( meta_context_var[idx], tf.cond(meta_period, lambda: meta_agent.context_vars[idx], lambda: meta_context_var[idx])) for idx in range(len(meta_context))] return tf.group( collect_experience_op, states_var_upd, state_var_upd, reward_var_upd, meta_action_var_upd, *meta_context_var_upd) def sample_best_meta_actions(state_reprs, next_state_reprs, prev_meta_actions, low_states, low_actions, low_state_reprs, inverse_dynamics, uvf_agent, k=10): """Return meta-actions which approximately maximize low-level log-probs.""" sampled_actions = inverse_dynamics.sample(state_reprs, next_state_reprs, k, prev_meta_actions) sampled_actions = tf.stop_gradient(sampled_actions) sampled_log_probs = tf.reshape(uvf_agent.log_probs( tf.tile(low_states, [k, 1, 1]), tf.tile(low_actions, [k, 1, 1]), tf.tile(low_state_reprs, [k, 1, 1]), [tf.reshape(sampled_actions, [-1, sampled_actions.shape[-1]])]), [k, low_states.shape[0], low_states.shape[1], -1]) fitness = tf.reduce_sum(sampled_log_probs, [2, 3]) best_actions = tf.argmax(fitness, 0) actions = tf.gather_nd( sampled_actions, tf.stack([best_actions, tf.range(prev_meta_actions.shape[0], dtype=tf.int64)], -1)) return actions @gin.configurable def train_uvf(train_dir, environment=None, num_bin_actions=3, agent_class=None, meta_agent_class=None, state_preprocess_class=None, inverse_dynamics_class=None, exp_action_wrapper=None, replay_buffer=None, meta_replay_buffer=None, replay_num_steps=1, meta_replay_num_steps=1, critic_optimizer=None, actor_optimizer=None, meta_critic_optimizer=None, meta_actor_optimizer=None, repr_optimizer=None, relabel_contexts=False, meta_relabel_contexts=False, batch_size=64, repeat_size=0, num_episodes_train=2000, initial_episodes=2, initial_steps=None, num_updates_per_observation=1, num_collect_per_update=1, num_collect_per_meta_update=1, gamma=1.0, meta_gamma=1.0, reward_scale_factor=1.0, target_update_period=1, should_stop_early=None, clip_gradient_norm=0.0, summarize_gradients=False, debug_summaries=False, log_every_n_steps=100, prefetch_queue_capacity=2, policy_save_dir='policy', save_policy_every_n_steps=1000, save_policy_interval_secs=0, replay_context_ratio=0.0, next_state_as_context_ratio=0.0, state_index=0, zero_timer_ratio=0.0, timer_index=-1, debug=False, max_policies_to_save=None, max_steps_per_episode=None, load_path=LOAD_PATH): """Train an agent.""" tf_env = create_maze_env.TFPyEnvironment(environment) observation_spec = [tf_env.observation_spec()] action_spec = [tf_env.action_spec()] max_steps_per_episode = max_steps_per_episode or tf_env.pyenv.max_episode_steps assert max_steps_per_episode, 'max_steps_per_episode need to be set' if initial_steps is None: initial_steps = initial_episodes * max_steps_per_episode if agent_class.ACTION_TYPE == 'discrete': assert False else: assert agent_class.ACTION_TYPE == 'continuous' assert agent_class.ACTION_TYPE == meta_agent_class.ACTION_TYPE with tf.variable_scope('meta_agent'): meta_agent = meta_agent_class( observation_spec, action_spec, tf_env, debug_summaries=debug_summaries) meta_agent.set_replay(replay=meta_replay_buffer) with tf.variable_scope('uvf_agent'): uvf_agent = agent_class( observation_spec, action_spec, tf_env, debug_summaries=debug_summaries) uvf_agent.set_meta_agent(agent=meta_agent) uvf_agent.set_replay(replay=replay_buffer) with tf.variable_scope('state_preprocess'): state_preprocess = state_preprocess_class() with tf.variable_scope('inverse_dynamics'): inverse_dynamics = inverse_dynamics_class( meta_agent.sub_context_as_action_specs[0]) # Create counter variables global_step = tf.contrib.framework.get_or_create_global_step() num_episodes = tf.Variable(0, dtype=tf.int64, name='num_episodes') num_resets = tf.Variable(0, dtype=tf.int64, name='num_resets') num_updates = tf.Variable(0, dtype=tf.int64, name='num_updates') num_meta_updates = tf.Variable(0, dtype=tf.int64, name='num_meta_updates') episode_rewards = tf.Variable([0.] * 100, name='episode_rewards') episode_meta_rewards = tf.Variable([0.] * 100, name='episode_meta_rewards') # Create counter variables summaries train_utils.create_counter_summaries([ ('environment_steps', global_step), ('num_episodes', num_episodes), ('num_resets', num_resets), ('num_updates', num_updates), ('num_meta_updates', num_meta_updates), ('replay_buffer_adds', replay_buffer.get_num_adds()), ('meta_replay_buffer_adds', meta_replay_buffer.get_num_adds()), ]) tf.summary.scalar('avg_episode_rewards', tf.reduce_mean(episode_rewards[1:])) tf.summary.scalar('avg_episode_meta_rewards', tf.reduce_mean(episode_meta_rewards[1:])) tf.summary.histogram('episode_rewards', episode_rewards[1:]) tf.summary.histogram('episode_meta_rewards', episode_meta_rewards[1:]) # Create init ops action_fn = uvf_agent.action action_fn = uvf_agent.add_noise_fn(action_fn, global_step=None) meta_action_fn = meta_agent.action meta_action_fn = meta_agent.add_noise_fn(meta_action_fn, global_step=None) meta_actions_fn = meta_agent.actions meta_actions_fn = meta_agent.add_noise_fn(meta_actions_fn, global_step=None) init_collect_experience_op = collect_experience( tf_env, uvf_agent, meta_agent, state_preprocess, replay_buffer, meta_replay_buffer, action_fn, meta_action_fn, environment_steps=global_step, num_episodes=num_episodes, num_resets=num_resets, episode_rewards=episode_rewards, episode_meta_rewards=episode_meta_rewards, store_context=True, disable_agent_reset=False, ) # Create train ops collect_experience_op = collect_experience( tf_env, uvf_agent, meta_agent, state_preprocess, replay_buffer, meta_replay_buffer, action_fn, meta_action_fn, environment_steps=global_step, num_episodes=num_episodes, num_resets=num_resets, episode_rewards=episode_rewards, episode_meta_rewards=episode_meta_rewards, store_context=True, disable_agent_reset=False, ) train_op_list = [] repr_train_op = tf.constant(0.0) for mode in ['meta', 'nometa']: if mode == 'meta': agent = meta_agent buff = meta_replay_buffer critic_opt = meta_critic_optimizer actor_opt = meta_actor_optimizer relabel = meta_relabel_contexts num_steps = meta_replay_num_steps my_gamma = meta_gamma, n_updates = num_meta_updates else: agent = uvf_agent buff = replay_buffer critic_opt = critic_optimizer actor_opt = actor_optimizer relabel = relabel_contexts num_steps = replay_num_steps my_gamma = gamma n_updates = num_updates with tf.name_scope(mode): batch = buff.get_random_batch(batch_size, num_steps=num_steps) states, actions, rewards, discounts, next_states = batch[:5] with tf.name_scope('Reward'): tf.summary.scalar('average_step_reward', tf.reduce_mean(rewards)) rewards *= reward_scale_factor batch_queue = slim.prefetch_queue.prefetch_queue( [states, actions, rewards, discounts, next_states] + batch[5:], capacity=prefetch_queue_capacity, name='batch_queue') batch_dequeue = batch_queue.dequeue() if repeat_size > 0: batch_dequeue = [ tf.tile(batch, (repeat_size+1,) + (1,) * (batch.shape.ndims - 1)) for batch in batch_dequeue ] batch_size *= (repeat_size + 1) states, actions, rewards, discounts, next_states = batch_dequeue[:5] if mode == 'meta': low_states = batch_dequeue[5] low_actions = batch_dequeue[6] low_state_reprs = state_preprocess(low_states) state_reprs = state_preprocess(states) next_state_reprs = state_preprocess(next_states) if mode == 'meta': # Re-label meta-action prev_actions = actions if FLAGS.goal_sample_strategy == 'None': pass elif FLAGS.goal_sample_strategy == 'FuN': actions = inverse_dynamics.sample(state_reprs, next_state_reprs, 1, prev_actions, sc=0.1) actions = tf.stop_gradient(actions) elif FLAGS.goal_sample_strategy == 'sample': actions = sample_best_meta_actions(state_reprs, next_state_reprs, prev_actions, low_states, low_actions, low_state_reprs, inverse_dynamics, uvf_agent, k=10) else: assert False if state_preprocess.trainable and mode == 'meta': # Representation learning is based on meta-transitions, but is trained # along with low-level policy updates. repr_loss, _, _ = state_preprocess.loss(states, next_states, low_actions, low_states) repr_train_op = slim.learning.create_train_op( repr_loss, repr_optimizer, global_step=None, update_ops=None, summarize_gradients=summarize_gradients, clip_gradient_norm=clip_gradient_norm, variables_to_train=state_preprocess.get_trainable_vars(),) # Get contexts for training contexts, next_contexts = agent.sample_contexts( mode='train', batch_size=batch_size, state=states, next_state=next_states, ) if not relabel: # Re-label context (in the style of TDM or HER). contexts, next_contexts = ( batch_dequeue[-2*len(contexts):-1*len(contexts)], batch_dequeue[-1*len(contexts):]) merged_states = agent.merged_states(states, contexts) merged_next_states = agent.merged_states(next_states, next_contexts) if mode == 'nometa': context_rewards, context_discounts = agent.compute_rewards( 'train', state_reprs, actions, rewards, next_state_reprs, contexts) elif mode == 'meta': # Meta-agent uses sum of rewards, not context-specific rewards. _, context_discounts = agent.compute_rewards( 'train', states, actions, rewards, next_states, contexts) context_rewards = rewards if agent.gamma_index is not None: context_discounts *= tf.cast( tf.reshape(contexts[agent.gamma_index], (-1,)), dtype=context_discounts.dtype) else: context_discounts *= my_gamma critic_loss = agent.critic_loss(merged_states, actions, context_rewards, context_discounts, merged_next_states) critic_loss = tf.reduce_mean(critic_loss) actor_loss = agent.actor_loss(merged_states, actions, context_rewards, context_discounts, merged_next_states) actor_loss *= tf.to_float( # Only update actor every N steps. tf.equal(n_updates % target_update_period, 0)) critic_train_op = slim.learning.create_train_op( critic_loss, critic_opt, global_step=n_updates, update_ops=None, summarize_gradients=summarize_gradients, clip_gradient_norm=clip_gradient_norm, variables_to_train=agent.get_trainable_critic_vars(),) critic_train_op = uvf_utils.tf_print( critic_train_op, [critic_train_op], message='critic_loss', print_freq=1000, name='critic_loss') train_op_list.append(critic_train_op) if actor_loss is not None: actor_train_op = slim.learning.create_train_op( actor_loss, actor_opt, global_step=None, update_ops=None, summarize_gradients=summarize_gradients, clip_gradient_norm=clip_gradient_norm, variables_to_train=agent.get_trainable_actor_vars(),) actor_train_op = uvf_utils.tf_print( actor_train_op, [actor_train_op], message='actor_loss', print_freq=1000, name='actor_loss') train_op_list.append(actor_train_op) assert len(train_op_list) == 4 # Update targets should happen after the networks have been updated. with tf.control_dependencies(train_op_list[2:]): update_targets_op = uvf_utils.periodically( uvf_agent.update_targets, target_update_period, 'update_targets') if meta_agent is not None: with tf.control_dependencies(train_op_list[:2]): update_meta_targets_op = uvf_utils.periodically( meta_agent.update_targets, target_update_period, 'update_targets') assert_op = tf.Assert( # Hack to get training to stop. tf.less_equal(global_step, 200 + num_episodes_train * max_steps_per_episode), [global_step]) with tf.control_dependencies([update_targets_op, assert_op]): train_op = tf.add_n(train_op_list[2:], name='post_update_targets') # Representation training steps on every low-level policy training step. train_op += repr_train_op with tf.control_dependencies([update_meta_targets_op, assert_op]): meta_train_op = tf.add_n(train_op_list[:2], name='post_update_meta_targets') if debug_summaries: train_.gen_debug_batch_summaries(batch) slim.summaries.add_histogram_summaries( uvf_agent.get_trainable_critic_vars(), 'critic_vars') slim.summaries.add_histogram_summaries( uvf_agent.get_trainable_actor_vars(), 'actor_vars') train_ops = train_utils.TrainOps(train_op, meta_train_op, collect_experience_op) policy_save_path = os.path.join(train_dir, policy_save_dir, 'model.ckpt') policy_vars = uvf_agent.get_actor_vars() + meta_agent.get_actor_vars() + [ global_step, num_episodes, num_resets ] + list(uvf_agent.context_vars) + list(meta_agent.context_vars) + state_preprocess.get_trainable_vars() # add critic vars, since some test evaluation depends on them policy_vars += uvf_agent.get_trainable_critic_vars() + meta_agent.get_trainable_critic_vars() policy_saver = tf.train.Saver( policy_vars, max_to_keep=max_policies_to_save, sharded=False) lowlevel_vars = (uvf_agent.get_actor_vars() + uvf_agent.get_trainable_critic_vars() + state_preprocess.get_trainable_vars()) lowlevel_saver = tf.train.Saver(lowlevel_vars) def policy_save_fn(sess): policy_saver.save( sess, policy_save_path, global_step=global_step, write_meta_graph=False) if save_policy_interval_secs > 0: tf.logging.info( 'Wait %d secs after save policy.' % save_policy_interval_secs) time.sleep(save_policy_interval_secs) train_step_fn = train_utils.TrainStep( max_number_of_steps=num_episodes_train * max_steps_per_episode + 100, num_updates_per_observation=num_updates_per_observation, num_collect_per_update=num_collect_per_update, num_collect_per_meta_update=num_collect_per_meta_update, log_every_n_steps=log_every_n_steps, policy_save_fn=policy_save_fn, save_policy_every_n_steps=save_policy_every_n_steps, should_stop_early=should_stop_early).train_step local_init_op = tf.local_variables_initializer() init_targets_op = tf.group(uvf_agent.update_targets(1.0), meta_agent.update_targets(1.0)) def initialize_training_fn(sess): """Initialize training function.""" sess.run(local_init_op) sess.run(init_targets_op) if load_path: tf.logging.info('Restoring low-level from %s' % load_path) lowlevel_saver.restore(sess, load_path) global_step_value = sess.run(global_step) assert global_step_value == 0, 'Global step should be zero.' collect_experience_call = sess.make_callable( init_collect_experience_op) for _ in range(initial_steps): collect_experience_call() train_saver = tf.train.Saver(max_to_keep=2, sharded=True) tf.logging.info('train dir: %s', train_dir) return slim.learning.train( train_ops, train_dir, train_step_fn=train_step_fn, save_interval_secs=FLAGS.save_interval_secs, saver=train_saver, log_every_n_steps=0, global_step=global_step, master="", is_chief=(FLAGS.task == 0), save_summaries_secs=FLAGS.save_summaries_secs, init_fn=initialize_training_fn)