# coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# 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"""Evaluation script for RL agents.

Example invocation:

python -m tensor2tensor.rl.evaluator \
    --policy_dir=$HOME/t2t/rl_v1/policy \
    --eval_metrics_dir=$HOME/t2t/rl_v1/full_eval_metrics \
    --hparams_set=rlmb_base \
    --hparams='batch_size=64'
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime
import os

from tensor2tensor.data_generators import gym_env
from tensor2tensor.layers import common_video
from tensor2tensor.models.research import rl  # pylint: disable=unused-import
from tensor2tensor.rl import rl_utils
from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib

import tensorflow as tf


flags = tf.flags
FLAGS = flags.FLAGS

flags.DEFINE_string("output_dir", "", "Main directory for multi-runs.")
flags.DEFINE_integer("total_num_workers", 1, "How many workers in total.")
flags.DEFINE_string("worker_to_game_map", "", "How to map workers to games.")
flags.DEFINE_string("policy_dir", "", "Directory with policy checkpoints.")
flags.DEFINE_string("model_dir", "", "Directory with model checkpoints.")
flags.DEFINE_string(
    "eval_metrics_dir", "", "Directory to output the eval metrics at."
)
flags.DEFINE_bool("full_eval", True, "Whether to ignore the timestep limit.")
flags.DEFINE_enum(
    "agent", "policy", ["random", "policy", "planner"], "Agent type to use."
)
flags.DEFINE_bool(
    "eval_with_learner", True,
    "Whether to use the PolicyLearner.evaluate function instead of an "
    "out-of-graph one. Works only with --agent=policy."
)
flags.DEFINE_string(
    "planner_hparams_set", "planner_small", "Planner hparam set."
)
flags.DEFINE_string("planner_hparams", "", "Planner hparam overrides.")
flags.DEFINE_integer(
    "log_every_steps", 20, "Log every how many environment steps."
)
flags.DEFINE_string(
    "debug_video_path", "", "Path to save the planner debug video at."
)

# Unused flags needed to pass for multi-run infrastructure.
flags.DEFINE_bool("autotune", False, "Unused here.")
flags.DEFINE_string("objective", "", "Unused here.")
flags.DEFINE_string("client_handle", "client_0", "Unused.")
flags.DEFINE_bool("maximize_tuner_objective", True, "Unused.")
flags.DEFINE_integer("vizier_search_algorithm", 0, "Unused.")


@registry.register_hparams
def planner_tiny():
  return tf.contrib.training.HParams(
      num_rollouts=1,
      planning_horizon=2,
      rollout_agent_type="random",
      batch_size=1,
      env_type="simulated",
  )


@registry.register_hparams
def planner_small():
  return tf.contrib.training.HParams(
      num_rollouts=64,
      planning_horizon=16,
      rollout_agent_type="policy",
      batch_size=64,
      env_type="simulated",
  )


def make_env(env_type, real_env, sim_env_kwargs):
  """Factory function for envs."""
  return {
      "real": lambda: real_env.new_like(  # pylint: disable=g-long-lambda
          batch_size=sim_env_kwargs["batch_size"],
          store_rollouts=False,
      ),
      "simulated": lambda: rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames(  # pylint: disable=g-long-lambda
          **sim_env_kwargs
      ),
  }[env_type]()


def make_agent(
    agent_type, env, policy_hparams, policy_dir, sampling_temp,
    sim_env_kwargs=None, frame_stack_size=None, planning_horizon=None,
    rollout_agent_type=None, batch_size=None, num_rollouts=None,
    inner_batch_size=None, video_writer=None, env_type=None):
  """Factory function for Agents."""
  if batch_size is None:
    batch_size = env.batch_size
  return {
      "random": lambda: rl_utils.RandomAgent(  # pylint: disable=g-long-lambda
          batch_size, env.observation_space, env.action_space
      ),
      "policy": lambda: rl_utils.PolicyAgent(  # pylint: disable=g-long-lambda
          batch_size, env.observation_space, env.action_space,
          policy_hparams, policy_dir, sampling_temp
      ),
      "planner": lambda: rl_utils.PlannerAgent(  # pylint: disable=g-long-lambda
          batch_size, make_agent(
              rollout_agent_type, env, policy_hparams, policy_dir,
              sampling_temp, batch_size=inner_batch_size
          ), make_env(env_type, env.env, sim_env_kwargs),
          lambda env: rl_utils.BatchStackWrapper(env, frame_stack_size),
          num_rollouts, planning_horizon,
          discount_factor=policy_hparams.gae_gamma, video_writer=video_writer
      ),
  }[agent_type]()


def make_eval_fn_with_agent(
    agent_type, planner_hparams, model_dir, log_every_steps=None,
    video_writer=None
):
  """Returns an out-of-graph eval_fn using the Agent API."""
  def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
    """Eval function."""
    base_env = env
    env = rl_utils.BatchStackWrapper(env, loop_hparams.frame_stack_size)
    sim_env_kwargs = rl.make_simulated_env_kwargs(
        base_env, loop_hparams, batch_size=planner_hparams.batch_size,
        model_dir=model_dir
    )
    agent = make_agent(
        agent_type, env, policy_hparams, policy_dir, sampling_temp,
        sim_env_kwargs, loop_hparams.frame_stack_size,
        planner_hparams.planning_horizon, planner_hparams.rollout_agent_type,
        num_rollouts=planner_hparams.num_rollouts,
        inner_batch_size=planner_hparams.batch_size, video_writer=video_writer,
        env_type=planner_hparams.env_type
    )
    rl_utils.run_rollouts(
        env, agent, env.reset(), log_every_steps=log_every_steps
    )
    assert len(base_env.current_epoch_rollouts()) == env.batch_size
  return eval_fn


def evaluate(
    loop_hparams, planner_hparams, policy_dir, model_dir, eval_metrics_dir,
    agent_type, eval_with_learner, log_every_steps, debug_video_path,
    report_fn=None, report_metric=None
):
  """Evaluate."""
  if eval_with_learner:
    assert agent_type == "policy"

  if report_fn:
    assert report_metric is not None

  eval_metrics_writer = tf.summary.FileWriter(eval_metrics_dir)
  video_writer = None
  kwargs = {}
  if not eval_with_learner:
    if debug_video_path:
      video_writer = common_video.WholeVideoWriter(
          fps=10, output_path=debug_video_path, file_format="avi")
    kwargs["eval_fn"] = make_eval_fn_with_agent(
        agent_type, planner_hparams, model_dir, log_every_steps=log_every_steps,
        video_writer=video_writer
    )
  eval_metrics = rl_utils.evaluate_all_configs(
      loop_hparams, policy_dir, **kwargs
  )
  rl_utils.summarize_metrics(eval_metrics_writer, eval_metrics, 0)

  if video_writer is not None:
    video_writer.finish_to_disk()

  # Report metrics
  if report_fn:
    if report_metric == "mean_reward":
      metric_name = rl_utils.get_metric_name(
          sampling_temp=loop_hparams.eval_sampling_temps[0],
          max_num_noops=loop_hparams.eval_max_num_noops,
          clipped=False
      )
      report_fn(eval_metrics[metric_name], 0)
    else:
      report_fn(eval_metrics[report_metric], 0)
  return eval_metrics


def get_game_for_worker(map_name, directory_id):
  """Get game for the given worker (directory) id."""
  if map_name == "v100unfriendly":
    games = ["chopper_command", "boxing", "asterix", "seaquest"]
    worker_per_game = 5
  elif map_name == "human_nice":
    games = gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE
    worker_per_game = 5
  else:
    raise ValueError("Unknown worker to game map name: %s" % map_name)
  games.sort()
  game_id = (directory_id - 1) // worker_per_game
  tf.logging.info("Getting game %d from %s." % (game_id, games))
  return games[game_id]


def main(_):
  now = datetime.datetime.now()
  now_tag = now.strftime("%Y_%m_%d_%H_%M")
  loop_hparams = trainer_lib.create_hparams(
      FLAGS.loop_hparams_set, FLAGS.loop_hparams
  )
  if FLAGS.worker_to_game_map and FLAGS.total_num_workers > 1:
    loop_hparams.game = get_game_for_worker(
        FLAGS.worker_to_game_map, FLAGS.worker_id + 1)
    tf.logging.info("Set game to %s." % loop_hparams.game)
  if FLAGS.full_eval:
    loop_hparams.eval_rl_env_max_episode_steps = -1
  planner_hparams = trainer_lib.create_hparams(
      FLAGS.planner_hparams_set, FLAGS.planner_hparams
  )
  policy_dir = FLAGS.policy_dir
  model_dir = FLAGS.model_dir
  eval_metrics_dir = FLAGS.eval_metrics_dir
  if FLAGS.output_dir:
    cur_dir = FLAGS.output_dir
    if FLAGS.total_num_workers > 1:
      cur_dir = os.path.join(cur_dir, "%d" % (FLAGS.worker_id + 1))
    policy_dir = os.path.join(cur_dir, "policy")
    model_dir = os.path.join(cur_dir, "world_model")
    eval_metrics_dir = os.path.join(cur_dir, "evaluator_" + now_tag)
    tf.logging.info("Writing metrics to %s." % eval_metrics_dir)
    if not tf.gfile.Exists(eval_metrics_dir):
      tf.gfile.MkDir(eval_metrics_dir)
  evaluate(
      loop_hparams, planner_hparams, policy_dir, model_dir,
      eval_metrics_dir, FLAGS.agent, FLAGS.eval_with_learner,
      FLAGS.log_every_steps if FLAGS.log_every_steps > 0 else None,
      debug_video_path=FLAGS.debug_video_path
  )


if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.app.run()

[((87, 9, 93, 3), 'tensorflow.contrib.training.HParams', 'tf.contrib.training.HParams', (), '', True, 'import tensorflow as tf\n'), ((98, 9, 104, 3), 'tensorflow.contrib.training.HParams', 'tf.contrib.training.HParams', (), '', True, 'import tensorflow as tf\n'), ((188, 24, 188, 63), 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', ({(188, 46, 188, 62): 'eval_metrics_dir'}, {}), '(eval_metrics_dir)', True, 'import tensorflow as tf\n'), ((199, 17, 201, 3), 'tensor2tensor.rl.rl_utils.evaluate_all_configs', 'rl_utils.evaluate_all_configs', ({(200, 6, 200, 18): 'loop_hparams', (200, 20, 200, 30): 'policy_dir'}, {}), '(loop_hparams, policy_dir, **kwargs)', False, 'from tensor2tensor.rl import rl_utils\n'), ((202, 2, 202, 66), 'tensor2tensor.rl.rl_utils.summarize_metrics', 'rl_utils.summarize_metrics', ({(202, 29, 202, 48): 'eval_metrics_writer', (202, 50, 202, 62): 'eval_metrics', (202, 64, 202, 65): '(0)'}, {}), '(eval_metrics_writer, eval_metrics, 0)', False, 'from tensor2tensor.rl import rl_utils\n'), ((233, 2, 233, 64), 'tensorflow.logging.info', 'tf.logging.info', ({(233, 18, 233, 63): "('Getting game %d from %s.' % (game_id, games))"}, {}), "('Getting game %d from %s.' % (game_id, games))", True, 'import tensorflow as tf\n'), ((238, 8, 238, 31), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((240, 17, 242, 3), 'tensor2tensor.utils.trainer_lib.create_hparams', 'trainer_lib.create_hparams', ({(241, 6, 241, 28): 'FLAGS.loop_hparams_set', (241, 30, 241, 48): 'FLAGS.loop_hparams'}, {}), '(FLAGS.loop_hparams_set, FLAGS.loop_hparams)', False, 'from tensor2tensor.utils import trainer_lib\n'), ((249, 20, 251, 3), 'tensor2tensor.utils.trainer_lib.create_hparams', 'trainer_lib.create_hparams', ({(250, 6, 250, 31): 'FLAGS.planner_hparams_set', (250, 33, 250, 54): 'FLAGS.planner_hparams'}, {}), '(FLAGS.planner_hparams_set, FLAGS.planner_hparams)', False, 'from tensor2tensor.utils import trainer_lib\n'), ((274, 2, 274, 43), 'tensorflow.logging.set_verbosity', 'tf.logging.set_verbosity', ({(274, 27, 274, 42): 'tf.logging.INFO'}, {}), '(tf.logging.INFO)', True, 'import tensorflow as tf\n'), ((275, 2, 275, 14), 'tensorflow.app.run', 'tf.app.run', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((156, 10, 156, 72), 'tensor2tensor.rl.rl_utils.BatchStackWrapper', 'rl_utils.BatchStackWrapper', ({(156, 37, 156, 40): 'env', (156, 42, 156, 71): 'loop_hparams.frame_stack_size'}, {}), '(env, loop_hparams.frame_stack_size)', False, 'from tensor2tensor.rl import rl_utils\n'), ((157, 21, 160, 5), 'tensor2tensor.models.research.rl.make_simulated_env_kwargs', 'rl.make_simulated_env_kwargs', (), '', False, 'from tensor2tensor.models.research import rl\n'), ((246, 4, 246, 58), 'tensorflow.logging.info', 'tf.logging.info', ({(246, 20, 246, 57): "('Set game to %s.' % loop_hparams.game)"}, {}), "('Set game to %s.' % loop_hparams.game)", True, 'import tensorflow as tf\n'), ((259, 17, 259, 48), 'os.path.join', 'os.path.join', ({(259, 30, 259, 37): 'cur_dir', (259, 39, 259, 47): '"""policy"""'}, {}), "(cur_dir, 'policy')", False, 'import os\n'), ((260, 16, 260, 52), 'os.path.join', 'os.path.join', ({(260, 29, 260, 36): 'cur_dir', (260, 38, 260, 51): '"""world_model"""'}, {}), "(cur_dir, 'world_model')", False, 'import os\n'), ((261, 23, 261, 68), 'os.path.join', 'os.path.join', ({(261, 36, 261, 43): 'cur_dir', (261, 45, 261, 67): "'evaluator_' + now_tag"}, {}), "(cur_dir, 'evaluator_' + now_tag)", False, 'import os\n'), ((262, 4, 262, 64), 'tensorflow.logging.info', 'tf.logging.info', ({(262, 20, 262, 63): "('Writing metrics to %s.' % eval_metrics_dir)"}, {}), "('Writing metrics to %s.' % eval_metrics_dir)", True, 'import tensorflow as tf\n'), ((193, 21, 194, 66), 'tensor2tensor.layers.common_video.WholeVideoWriter', 'common_video.WholeVideoWriter', (), '', False, 'from tensor2tensor.layers import common_video\n'), ((210, 20, 214, 7), 'tensor2tensor.rl.rl_utils.get_metric_name', 'rl_utils.get_metric_name', (), '', False, 'from tensor2tensor.rl import rl_utils\n'), ((258, 16, 258, 67), 'os.path.join', 'os.path.join', ({(258, 29, 258, 36): 'cur_dir', (258, 38, 258, 66): "'%d' % (FLAGS.worker_id + 1)"}, {}), "(cur_dir, '%d' % (FLAGS.worker_id + 1))", False, 'import os\n'), ((263, 11, 263, 44), 'tensorflow.gfile.Exists', 'tf.gfile.Exists', ({(263, 27, 263, 43): 'eval_metrics_dir'}, {}), '(eval_metrics_dir)', True, 'import tensorflow as tf\n'), ((264, 6, 264, 38), 'tensorflow.gfile.MkDir', 'tf.gfile.MkDir', ({(264, 21, 264, 37): 'eval_metrics_dir'}, {}), '(eval_metrics_dir)', True, 'import tensorflow as tf\n'), ((114, 27, 116, 7), 'tensor2tensor.rl.rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames', 'rl_utils.SimulatedBatchGymEnvWithFixedInitialFrames', ({}, {}), '(**sim_env_kwargs)', False, 'from tensor2tensor.rl import rl_utils\n'), ((129, 24, 131, 7), 'tensor2tensor.rl.rl_utils.RandomAgent', 'rl_utils.RandomAgent', ({(130, 10, 130, 20): 'batch_size', (130, 22, 130, 43): 'env.observation_space', (130, 45, 130, 61): 'env.action_space'}, {}), '(batch_size, env.observation_space, env.action_space)', False, 'from tensor2tensor.rl import rl_utils\n'), ((132, 24, 135, 7), 'tensor2tensor.rl.rl_utils.PolicyAgent', 'rl_utils.PolicyAgent', ({(133, 10, 133, 20): 'batch_size', (133, 22, 133, 43): 'env.observation_space', (133, 45, 133, 61): 'env.action_space', (134, 10, 134, 24): 'policy_hparams', (134, 26, 134, 36): 'policy_dir', (134, 38, 134, 51): 'sampling_temp'}, {}), '(batch_size, env.observation_space, env.action_space,\n    policy_hparams, policy_dir, sampling_temp)', False, 'from tensor2tensor.rl import rl_utils\n'), ((141, 22, 141, 71), 'tensor2tensor.rl.rl_utils.BatchStackWrapper', 'rl_utils.BatchStackWrapper', ({(141, 49, 141, 52): 'env', (141, 54, 141, 70): 'frame_stack_size'}, {}), '(env, frame_stack_size)', False, 'from tensor2tensor.rl import rl_utils\n')]