text: #!/usr/bin/env python
  
  """
  Run a trained checkpoint to see what the agent is actually doing in the
  environment.
  """
  
  import argparse
  import os.path as osp
  import time
  from collections import deque
  
  import cloudpickle
  import numpy as np
  import tensorflow as tf
  
  import matplotlib.pyplot as plt
  from matplotlib.ticker import FormatStrFormatter
  from utils import make_env
  
  
  def main():
      args = parse_args()
  
      env = make_env(args.env)
      model = get_model(args.policy_ckpt_dir)
      if args.reward_predictor_ckpt_dir:
          reward_predictor = get_reward_predictor(args.reward_predictor_ckpt_dir)
      else:
          reward_predictor = None
  
      run_agent(env, model, reward_predictor, args.frame_interval_ms)
  
  
  def run_agent(env, model, reward_predictor, frame_interval_ms):
      nenvs = 1
      nstack = int(model.step_model.X.shape[-1])
      nh, nw, nc = env.observation_space.shape
      obs = np.zeros((nenvs, nh, nw, nc * nstack), dtype=np.uint8)
      model_nenvs = int(model.step_model.X.shape[0])
      states = model.initial_state
      if reward_predictor:
          value_graph = ValueGraph()
      while True:
          raw_obs = env.reset()
          update_obs(obs, raw_obs, nc)
          episode_reward = 0
          done = False
          while not done:
              model_obs = np.vstack([obs] * model_nenvs)
              actions, _, states = model.step(model_obs, states, [done])
              action = actions[0]
              raw_obs, reward, done, _ = env.step(action)
              obs = update_obs(obs, raw_obs, nc)
              episode_reward += reward
              env.render()
              if reward_predictor is not None:
                  predicted_reward = reward_predictor.reward(obs)
                  # reward_predictor.reward returns reward for each frame in the
                  # supplied batch. We only supplied one frame, so get the reward
                  # for that frame.
                  value_graph.append(predicted_reward[0])
              time.sleep(frame_interval_ms * 1e-3)
          print("Episode reward:", episode_reward)
  
  
  def update_obs(obs, raw_obs, nc):
      obs = np.roll(obs, shift=-nc, axis=3)
      obs[:, :, :, -nc:] = raw_obs
      return obs
  
  
  def get_reward_predictor(ckpt_dir):
      with open(osp.join(ckpt_dir, 'make_reward_predictor.pkl'), 'rb') as fh:
          make_reward_predictor = cloudpickle.loads(fh.read())
      cluster_dict = {'a2c': ['localhost:2200']}
      print("Initialising reward predictor...")
      reward_predictor = make_reward_predictor(name='a2c', cluster_dict=cluster_dict)
      reward_predictor.init_network(ckpt_dir)
      return reward_predictor
  
  
  def get_model(ckpt_dir):
      model_file = osp.join(ckpt_dir, 'make_model.pkl')
      with open(model_file, 'rb') as fh:
          make_model = cloudpickle.loads(fh.read())
      print("Initialising policy...")
      model = make_model()
      ckpt_file = tf.train.latest_checkpoint(ckpt_dir)
      print("Loading checkpoint...")
      model.load(ckpt_file)
      return model
  
  
  def parse_args():
      parser = argparse.ArgumentParser()
      parser.add_argument("env")
      parser.add_argument("policy_ckpt_dir")
      parser.add_argument("--reward_predictor_ckpt_dir")
      parser.add_argument("--frame_interval_ms", type=float, default=0.)
      args = parser.parse_args()
      return args
  
  
  class ValueGraph:
      def __init__(self):
          n_values = 100
          self.data = deque(maxlen=n_values)
  
          self.fig, self.ax = plt.subplots()
          self.ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
          self.fig.set_size_inches(4, 2)
          self.ax.set_xlim([0, n_values - 1])
          self.ax.grid(axis='y')  # Draw a line at 0 reward
          self.y_min = float('inf')
          self.y_max = -float('inf')
          self.line, = self.ax.plot([], [])
  
          self.fig.show()
          self.fig.canvas.draw()
  
      def append(self, value):
          self.data.append(value)
  
          self.y_min = min(self.y_min, min(self.data))
          self.y_max = max(self.y_max, max(self.data))
          self.ax.set_ylim([self.y_min, self.y_max])
          self.ax.set_yticks([self.y_min, 0, self.y_max])
          plt.tight_layout()
  
          ydata = list(self.data)
          xdata = list(range(len(self.data)))
          self.line.set_data(xdata, ydata)
  
          self.ax.draw_artist(self.line)
          self.fig.canvas.draw()
  
  
  if __name__ == '__main__':
      main()
  