research/efficient-hrl/train.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"""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)