research/efficient-hrl/eval.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 evaluating a UVF agent.
To run locally: See run_eval.py
To run on borg: See train_eval.borg
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import tensorflow as tf
slim = tf.contrib.slim
import gin.tf
# pylint: disable=unused-import
import agent
import train
from utils import utils as uvf_utils
from utils import eval_utils
from environments import create_maze_env
# pylint: enable=unused-import
flags = tf.app.flags
flags.DEFINE_string('eval_dir', None,
'Directory for writing logs/summaries during eval.')
flags.DEFINE_string('checkpoint_dir', None,
'Directory containing checkpoints to eval.')
FLAGS = flags.FLAGS
def get_evaluate_checkpoint_fn(master, output_dir, eval_step_fns,
model_rollout_fn, gamma, max_steps_per_episode,
num_episodes_eval, num_episodes_videos,
tuner_hook, generate_videos,
generate_summaries, video_settings):
"""Returns a function that evaluates a given checkpoint.
Args:
master: BNS name of the TensorFlow master
output_dir: The output directory to which the metric summaries are written.
eval_step_fns: A dictionary of a functions that return a list of
[state, action, reward, discount, transition_type] tensors,
indexed by summary tag name.
model_rollout_fn: Model rollout fn.
gamma: Discount factor for the reward.
max_steps_per_episode: Maximum steps to run each episode for.
num_episodes_eval: Number of episodes to evaluate and average reward over.
num_episodes_videos: Number of episodes to record for video.
tuner_hook: A callable(average reward, global step) that updates a Vizier
tuner trial.
generate_videos: Whether to generate videos of the agent in action.
generate_summaries: Whether to generate summaries.
video_settings: Settings for generating videos of the agent.
Returns:
A function that evaluates a checkpoint.
"""
sess = tf.Session(master, graph=tf.get_default_graph())
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
summary_writer = tf.summary.FileWriter(output_dir)
def evaluate_checkpoint(checkpoint_path):
"""Performs a one-time evaluation of the given checkpoint.
Args:
checkpoint_path: Checkpoint to evaluate.
Returns:
True if the evaluation process should stop
"""
restore_fn = tf.contrib.framework.assign_from_checkpoint_fn(
checkpoint_path,
uvf_utils.get_all_vars(),
ignore_missing_vars=True,
reshape_variables=False)
assert restore_fn is not None, 'cannot restore %s' % checkpoint_path
restore_fn(sess)
global_step = sess.run(slim.get_global_step())
should_stop = False
max_reward = -1e10
max_meta_reward = -1e10
for eval_tag, (eval_step, env_base,) in sorted(eval_step_fns.items()):
if hasattr(env_base, 'set_sess'):
env_base.set_sess(sess) # set session
if generate_summaries:
tf.logging.info(
'[%s] Computing average reward over %d episodes at global step %d.',
eval_tag, num_episodes_eval, global_step)
(average_reward, last_reward,
average_meta_reward, last_meta_reward, average_success,
states, actions) = eval_utils.compute_average_reward(
sess, env_base, eval_step, gamma, max_steps_per_episode,
num_episodes_eval)
tf.logging.info('[%s] Average reward = %f', eval_tag, average_reward)
tf.logging.info('[%s] Last reward = %f', eval_tag, last_reward)
tf.logging.info('[%s] Average meta reward = %f', eval_tag, average_meta_reward)
tf.logging.info('[%s] Last meta reward = %f', eval_tag, last_meta_reward)
tf.logging.info('[%s] Average success = %f', eval_tag, average_success)
if model_rollout_fn is not None:
preds, model_losses = eval_utils.compute_model_loss(
sess, model_rollout_fn, states, actions)
for i, (pred, state, model_loss) in enumerate(
zip(preds, states, model_losses)):
tf.logging.info('[%s] Model rollout step %d: loss=%f', eval_tag, i,
model_loss)
tf.logging.info('[%s] Model rollout step %d: pred=%s', eval_tag, i,
str(pred.tolist()))
tf.logging.info('[%s] Model rollout step %d: state=%s', eval_tag, i,
str(state.tolist()))
# Report the eval stats to the tuner.
if average_reward > max_reward:
max_reward = average_reward
if average_meta_reward > max_meta_reward:
max_meta_reward = average_meta_reward
for (tag, value) in [('Reward/average_%s_reward', average_reward),
('Reward/last_%s_reward', last_reward),
('Reward/average_%s_meta_reward', average_meta_reward),
('Reward/last_%s_meta_reward', last_meta_reward),
('Reward/average_%s_success', average_success)]:
summary_str = tf.Summary(value=[
tf.Summary.Value(
tag=tag % eval_tag,
simple_value=value)
])
summary_writer.add_summary(summary_str, global_step)
summary_writer.flush()
if generate_videos or should_stop:
# Do a manual reset before generating the video to see the initial
# pose of the robot, towards which the reset controller is moving.
if hasattr(env_base, '_gym_env'):
tf.logging.info('Resetting before recording video')
if hasattr(env_base._gym_env, 'reset_model'):
env_base._gym_env.reset_model() # pylint: disable=protected-access
else:
env_base._gym_env.wrapped_env.reset_model()
video_filename = os.path.join(output_dir, 'videos',
'%s_step_%d.mp4' % (eval_tag,
global_step))
eval_utils.capture_video(sess, eval_step, env_base,
max_steps_per_episode * num_episodes_videos,
video_filename, video_settings,
reset_every=max_steps_per_episode)
should_stop = should_stop or (generate_summaries and tuner_hook and
tuner_hook(max_reward, global_step))
return bool(should_stop)
return evaluate_checkpoint
def get_model_rollout(uvf_agent, tf_env):
"""Model rollout function."""
state_spec = tf_env.observation_spec()[0]
action_spec = tf_env.action_spec()[0]
state_ph = tf.placeholder(dtype=state_spec.dtype, shape=state_spec.shape)
action_ph = tf.placeholder(dtype=action_spec.dtype, shape=action_spec.shape)
merged_state = uvf_agent.merged_state(state_ph)
diff_value = uvf_agent.critic_net(tf.expand_dims(merged_state, 0),
tf.expand_dims(action_ph, 0))[0]
diff_value = tf.cast(diff_value, dtype=state_ph.dtype)
state_ph.shape.assert_is_compatible_with(diff_value.shape)
next_state = state_ph + diff_value
def model_rollout_fn(sess, state, action):
return sess.run(next_state, feed_dict={state_ph: state, action_ph: action})
return model_rollout_fn
def get_eval_step(uvf_agent,
state_preprocess,
tf_env,
action_fn,
meta_action_fn,
environment_steps,
num_episodes,
mode='eval'):
"""Get one-step policy/env stepping ops.
Args:
uvf_agent: A UVF agent.
tf_env: A TFEnvironment.
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.
mode: a string representing the mode=[train, explore, eval].
Returns:
A collect_experience_op that excute an action and store into the
replay_buffer
"""
tf_env.start_collect()
state = tf_env.current_obs()
action = action_fn(state, context=None)
state_repr = state_preprocess(state)
action_spec = tf_env.action_spec()
action_ph = tf.placeholder(dtype=action_spec.dtype, shape=action_spec.shape)
with tf.control_dependencies([state]):
transition_type, reward, discount = tf_env.step(action_ph)
def increment_step():
return environment_steps.assign_add(1)
def increment_episode():
return num_episodes.assign_add(1)
def no_op_int():
return tf.constant(0, dtype=tf.int64)
step_cond = uvf_agent.step_cond_fn(state, action,
transition_type,
environment_steps, num_episodes)
reset_episode_cond = uvf_agent.reset_episode_cond_fn(
state, action,
transition_type, environment_steps, num_episodes)
reset_env_cond = uvf_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)
with tf.control_dependencies([increment_step_op]):
increment_episode_op = tf.cond(reset_episode_cond, increment_episode,
no_op_int)
with tf.control_dependencies([reward, discount]):
next_state = tf_env.current_obs()
next_state_repr = state_preprocess(next_state)
with tf.control_dependencies([increment_episode_op]):
post_reward, post_meta_reward = uvf_agent.cond_begin_episode_op(
tf.logical_not(reset_episode_cond),
[state, action_ph, reward, next_state,
state_repr, next_state_repr],
mode=mode, meta_action_fn=meta_action_fn)
# Important: do manual reset after getting the final reward from the
# unreset environment.
with tf.control_dependencies([post_reward, post_meta_reward]):
cond_reset_op = tf.cond(reset_env_cond,
tf_env.reset,
tf_env.current_time_step)
# Add a dummy control dependency to force the reset_op to run
with tf.control_dependencies(cond_reset_op):
post_reward, post_meta_reward = map(tf.identity, [post_reward, post_meta_reward])
eval_step = [next_state, action_ph, transition_type, post_reward, post_meta_reward, discount, uvf_agent.context_vars, state_repr]
if callable(action):
def step_fn(sess):
action_value = action(sess)
return sess.run(eval_step, feed_dict={action_ph: action_value})
else:
action = uvf_utils.clip_to_spec(action, action_spec)
def step_fn(sess):
action_value = sess.run(action)
return sess.run(eval_step, feed_dict={action_ph: action_value})
return step_fn
@gin.configurable
def evaluate(checkpoint_dir,
eval_dir,
environment=None,
num_bin_actions=3,
agent_class=None,
meta_agent_class=None,
state_preprocess_class=None,
gamma=1.0,
num_episodes_eval=10,
eval_interval_secs=60,
max_number_of_evaluations=None,
checkpoint_timeout=None,
timeout_fn=None,
tuner_hook=None,
generate_videos=False,
generate_summaries=True,
num_episodes_videos=5,
video_settings=None,
eval_modes=('eval',),
eval_model_rollout=False,
policy_save_dir='policy',
checkpoint_range=None,
checkpoint_path=None,
max_steps_per_episode=None,
evaluate_nohrl=False):
"""Loads and repeatedly evaluates a checkpointed model at a set interval.
Args:
checkpoint_dir: The directory where the checkpoints reside.
eval_dir: Directory to save the evaluation summary results.
environment: A BaseEnvironment to evaluate.
num_bin_actions: Number of bins for discretizing continuous actions.
agent_class: An RL agent class.
meta_agent_class: A Meta agent class.
gamma: Discount factor for the reward.
num_episodes_eval: Number of episodes to evaluate and average reward over.
eval_interval_secs: The number of seconds between each evaluation run.
max_number_of_evaluations: The max number of evaluations. If None the
evaluation continues indefinitely.
checkpoint_timeout: The maximum amount of time to wait between checkpoints.
If left as `None`, then the process will wait indefinitely.
timeout_fn: Optional function to call after a timeout.
tuner_hook: A callable that takes the average reward and global step and
updates a Vizier tuner trial.
generate_videos: Whether to generate videos of the agent in action.
generate_summaries: Whether to generate summaries.
num_episodes_videos: Number of episodes to evaluate for generating videos.
video_settings: Settings for generating videos of the agent.
optimal action based on the critic.
eval_modes: A tuple of eval modes.
eval_model_rollout: Evaluate model rollout.
policy_save_dir: Optional sub-directory where the policies are
saved.
checkpoint_range: Optional. If provided, evaluate all checkpoints in
the range.
checkpoint_path: Optional sub-directory specifying which checkpoint to
evaluate. If None, will evaluate the most recent checkpoint.
"""
tf_env = create_maze_env.TFPyEnvironment(environment)
observation_spec = [tf_env.observation_spec()]
action_spec = [tf_env.action_spec()]
assert max_steps_per_episode, 'max_steps_per_episode need to be set'
if agent_class.ACTION_TYPE == 'discrete':
assert False
else:
assert agent_class.ACTION_TYPE == 'continuous'
if meta_agent_class is not None:
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,
)
else:
meta_agent = None
with tf.variable_scope('uvf_agent'):
uvf_agent = agent_class(
observation_spec,
action_spec,
tf_env,
)
uvf_agent.set_meta_agent(agent=meta_agent)
with tf.variable_scope('state_preprocess'):
state_preprocess = state_preprocess_class()
# run both actor and critic once to ensure networks are initialized
# and gin configs will be saved
# pylint: disable=protected-access
temp_states = tf.expand_dims(
tf.zeros(
dtype=uvf_agent._observation_spec.dtype,
shape=uvf_agent._observation_spec.shape), 0)
# pylint: enable=protected-access
temp_actions = uvf_agent.actor_net(temp_states)
uvf_agent.critic_net(temp_states, temp_actions)
# create eval_step_fns for each action function
eval_step_fns = dict()
meta_agent = uvf_agent.meta_agent
for meta in [True] + [False] * evaluate_nohrl:
meta_tag = 'hrl' if meta else 'nohrl'
uvf_agent.set_meta_agent(meta_agent if meta else None)
for mode in eval_modes:
# wrap environment
wrapped_environment = uvf_agent.get_env_base_wrapper(
environment, mode=mode)
action_wrapper = lambda agent_: agent_.action
action_fn = action_wrapper(uvf_agent)
meta_action_fn = action_wrapper(meta_agent)
eval_step_fns['%s_%s' % (mode, meta_tag)] = (get_eval_step(
uvf_agent=uvf_agent,
state_preprocess=state_preprocess,
tf_env=tf_env,
action_fn=action_fn,
meta_action_fn=meta_action_fn,
environment_steps=tf.Variable(
0, dtype=tf.int64, name='environment_steps'),
num_episodes=tf.Variable(0, dtype=tf.int64, name='num_episodes'),
mode=mode), wrapped_environment,)
model_rollout_fn = None
if eval_model_rollout:
model_rollout_fn = get_model_rollout(uvf_agent, tf_env)
tf.train.get_or_create_global_step()
if policy_save_dir:
checkpoint_dir = os.path.join(checkpoint_dir, policy_save_dir)
tf.logging.info('Evaluating policies at %s', checkpoint_dir)
tf.logging.info('Running episodes for max %d steps', max_steps_per_episode)
evaluate_checkpoint_fn = get_evaluate_checkpoint_fn(
'', eval_dir, eval_step_fns, model_rollout_fn, gamma,
max_steps_per_episode, num_episodes_eval, num_episodes_videos, tuner_hook,
generate_videos, generate_summaries, video_settings)
if checkpoint_path is not None:
checkpoint_path = os.path.join(checkpoint_dir, checkpoint_path)
evaluate_checkpoint_fn(checkpoint_path)
elif checkpoint_range is not None:
model_files = tf.gfile.Glob(
os.path.join(checkpoint_dir, 'model.ckpt-*.index'))
tf.logging.info('Found %s policies at %s', len(model_files), checkpoint_dir)
model_files = {
int(f.split('model.ckpt-', 1)[1].split('.', 1)[0]):
os.path.splitext(f)[0]
for f in model_files
}
model_files = {
k: v
for k, v in model_files.items()
if k >= checkpoint_range[0] and k <= checkpoint_range[1]
}
tf.logging.info('Evaluating %d policies at %s',
len(model_files), checkpoint_dir)
for _, checkpoint_path in sorted(model_files.items()):
evaluate_checkpoint_fn(checkpoint_path)
else:
eval_utils.evaluate_checkpoint_repeatedly(
checkpoint_dir,
evaluate_checkpoint_fn,
eval_interval_secs=eval_interval_secs,
max_number_of_evaluations=max_number_of_evaluations,
checkpoint_timeout=checkpoint_timeout,
timeout_fn=timeout_fn)