# @title Install necessary packages.
!pip install -U dopamine-rl

# @title Necessary imports and globals.

import numpy as np
import os
from dopamine.agents.dqn import dqn_agent
from dopamine.discrete_domains import run_experiment
from dopamine.colab import utils as colab_utils
from absl import flags

BASE_PATH = '/tmp/colab_dope_run'  # @param
GAMES = ['Asterix', 'Pong', 'SpaceInvaders']  # @param

# @title Load the sample log files.

# For illustrative purposes, we are providing sample logs of the Rainbow agent
# trained without sticky actions.
!gsutil -q -m cp -R gs://download-dopamine-rl/colab/* /content/

"""
Explanation: Copyright 2018 The Dopamine 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
https://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.
Loading experiment statistics
This colab illustrates how you can load an experiment's statistics and plot it. We provide a sample statistics file for a
modified Rainbow agent trained on Dopamine. Note that the performance of this sample data is not reflective of the standard settings, it was compiled solely for illustrative purposes.

Example 1 shows how to use our colab utils to load the baselines and plot them against your experiment. 
Example 2 shows how to load the raw experiment statistics and plot using a different package.

To re-run this colab, run each cell in order.
End of explanation
"""


# @title Load the baseline data

!gsutil -q -m cp -R gs://download-dopamine-rl/preprocessed-benchmarks/* /content/
experimental_data = colab_utils.load_baselines('/content')

# @title Load the summarized sample data and add it to the global data frame.
import collections

# Now read the data for our samples in a summarized fashion. The files will be
# in the local directory /content/samples/rainbow/GAME_v4, so we make use of
# the parameter_set and job_descriptor parameters.
parameter_set = collections.OrderedDict([
    ('agent', ['rainbow']),
    ('game', GAMES)
])
sample_data = colab_utils.read_experiment(
    '/content/samples',
    parameter_set=parameter_set,
    job_descriptor='{}/{}_v4',
    summary_keys=['train_episode_returns'])
sample_data['agent'] = 'Sample Rainbow'
sample_data['run_number'] = 1
for game in GAMES:
  experimental_data[game] = experimental_data[game].merge(
      sample_data[sample_data.game == game], how='outer')

# @title Plot the sample agent data against the baselines.

import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt

for game in GAMES:
  fig, ax = plt.subplots(figsize=(16,8))
  sns.lineplot(x='iteration', y='train_episode_returns', hue='agent',
               data=experimental_data[game], ax=ax)
  plt.title(game)
  plt.show()

"""
Explanation: Example 1: Plot against Dopamine baselines
End of explanation
"""


import matplotlib.pyplot as plt

for game in GAMES:
  # Use our provided colab utils to load this log file. The second returned 
  raw_data, _ = colab_utils.load_statistics(
      '/content/samples/rainbow/{}_v4/logs'.format(game), verbose=False)
  summarized_data = colab_utils.summarize_data(
      raw_data, ['train_episode_returns'])
  plt.plot(summarized_data['train_episode_returns'], label='episode returns')
  plt.plot()
  plt.title('Rainbow training - {}'.format(game))
  plt.xlabel('Iteration')
  plt.ylabel('Return')
  plt.legend()
  plt.show()

"""
Explanation: Example 2: Load the raw data for our sample experiments and plot using a different plotting package.
End of explanation
"""


