class_docstr: Train and evaluate DQN on Atari.
class_docstr_tok: ["Train","and","evaluate","DQN","on","Atari","."]
docstr: A simple Atari train and eval for DQN.
  
  Args:
    root_dir: Directory to write log files to.
    env_name: Fully-qualified name of the Atari environment (i.e. Pong-v0).
    num_iterations: Number of train/eval iterations to run.
    max_episode_frames: Maximum length of a single episode, in ALE frames.
    terminal_on_life_loss: Whether to simulate an episode termination when a
      life is lost.
    conv_layer_params: Params for convolutional layers of QNetwork.
    fc_layer_params: Params for fully connected layers of QNetwork.
    initial_collect_steps: Number of frames to ALE frames to process before
      beginning to train. Since this is in ALE frames, there will be
      initial_collect_steps/4 items in the replay buffer when training starts.
    epsilon_greedy: Final epsilon value to decay to for training.
    epsilon_decay_period: Period over which to decay epsilon, from 1.0 to
      epsilon_greedy (defined above).
    replay_buffer_capacity: Maximum number of items to store in the replay
      buffer.
    train_steps_per_iteration: Number of ALE frames to run through for each
      iteration of training.
    update_period: Run a train operation every update_period ALE frames.
    target_update_tau: Coeffecient for soft target network updates (1.0 ==
      hard updates).
    target_update_period: Period, in ALE frames, to copy the live network to
      the target network.
    batch_size: Number of frames to include in each training batch.
    learning_rate: RMS optimizer learning rate.
    n_step_update: The number of steps to consider when computing TD error and
      TD loss. Applies standard single-step updates when set to 1.
    gamma: Discount for future rewards.
    reward_scale_factor: Scaling factor for rewards.
    gradient_clipping: Norm length to clip gradients.
    do_eval: If True, run an eval every iteration. If False, skip eval.
    eval_steps_per_iteration: Number of ALE frames to run through for each
      iteration of evaluation.
    eval_epsilon_greedy: Epsilon value to use for the evaluation policy (0 ==
      totally greedy policy).
    log_interval: Log stats to the terminal every log_interval training steps.
    summary_interval: Write TF summaries every summary_interval training
      steps.
    summaries_flush_secs: Flush summaries to disk every summaries_flush_secs
      seconds.
    debug_summaries: If True, write additional summaries for debugging (see
      dqn_agent for which summaries are written).
    summarize_grads_and_vars: Include gradients in summaries.
    eval_metrics_callback: A callback function that takes (metric_dict,
      global_step) as parameters. Called after every eval with the results of
      the evaluation.
docstr_tok: ["A","simple","Atari","train","and","eval","for","DQN",".","Args",":","root_dir",":","Directory","to","write","log","files","to",".","env_name",":","Fully-qualified","name","of","the","Atari","environment","(","i.e",".","Pong-v0",")",".","num_iterations",":","Number","of","train\/eval","iterations","to","run",".","max_episode_frames",":","Maximum","length","of","a","single","episode",",","in","ALE","frames",".","terminal_on_life_loss",":","Whether","to","simulate","an","episode","termination","when","a","life","is","lost",".","conv_layer_params",":","Params","for","convolutional","layers","of","QNetwork",".","fc_layer_params",":","Params","for","fully","connected","layers","of","QNetwork",".","initial_collect_steps",":","Number","of","frames","to","ALE","frames","to","process","before","beginning","to","train",".","Since","this","is","in","ALE","frames",",","there","will","be","initial_collect_steps\/4","items","in","the","replay","buffer","when","training","starts",".","epsilon_greedy",":","Final","epsilon","value","to","decay","to","for","training",".","epsilon_decay_period",":","Period","over","which","to","decay","epsilon",",","from","1.0","to","epsilon_greedy","(","defined","above",")",".","replay_buffer_capacity",":","Maximum","number","of","items","to","store","in","the","replay","buffer",".","train_steps_per_iteration",":","Number","of","ALE","frames","to","run","through","for","each","iteration","of","training",".","update_period",":","Run","a","train","operation","every","update_period","ALE","frames",".","target_update_tau",":","Coeffecient","for","soft","target","network","updates","(","1.0","==","hard","updates",")",".","target_update_period",":","Period",",","in","ALE","frames",",","to","copy","the","live","network","to","the","target","network",".","batch_size",":","Number","of","frames","to","include","in","each","training","batch",".","learning_rate",":","RMS","optimizer","learning","rate",".","n_step_update",":","The","number","of","steps","to","consider","when","computing","TD","error","and","TD","loss",".","Applies","standard","single-step","updates","when","set","to","1.","gamma",":","Discount","for","future","rewards",".","reward_scale_factor",":","Scaling","factor","for","rewards",".","gradient_clipping",":","Norm","length","to","clip","gradients",".","do_eval",":","If","True",",","run","an","eval","every","iteration",".","If","False",",","skip","eval",".","eval_steps_per_iteration",":","Number","of","ALE","frames","to","run","through","for","each","iteration","of","evaluation",".","eval_epsilon_greedy",":","Epsilon","value","to","use","for","the","evaluation","policy","(","0","==","totally","greedy","policy",")",".","log_interval",":","Log","stats","to","the","terminal","every","log_interval","training","steps",".","summary_interval",":","Write","TF","summaries","every","summary_interval","training","steps",".","summaries_flush_secs",":","Flush","summaries","to","disk","every","summaries_flush_secs","seconds",".","debug_summaries",":","If","True",",","write","additional","summaries","for","debugging","(","see","dqn_agent","for","which","summaries","are","written",")",".","summarize_grads_and_vars",":","Include","gradients","in","summaries",".","eval_metrics_callback",":","A","callback","function","that","takes","(","metric_dict",",","global_step",")","as","parameters",".","Called","after","every","eval","with","the","results","of","the","evaluation","."]
code: def __init__(
      self,
      root_dir,
      env_name,
      num_iterations=200,
      max_episode_frames=108000,  # ALE frames
      terminal_on_life_loss=False,
      conv_layer_params=(
          (32, (8, 8), 4),
          (64, (4, 4), 2),
          (64, (3, 3), 1),
      ),
      fc_layer_params=(512,),
      # Params for collect
      initial_collect_steps=80000,  # ALE frames
      epsilon_greedy=0.01,
      epsilon_decay_period=1000000,  # ALE frames
      replay_buffer_capacity=1000000,
      # Params for train
      train_steps_per_iteration=1000000,  # ALE frames
      update_period=16,  # ALE frames
      target_update_tau=1.0,
      target_update_period=32000,  # ALE frames
      batch_size=32,
      learning_rate=2.5e-4,
      n_step_update=2,
      gamma=0.99,
      reward_scale_factor=1.0,
      gradient_clipping=None,
      # Params for eval
      do_eval=True,
      eval_steps_per_iteration=500000,  # ALE frames
      eval_epsilon_greedy=0.001,
      # Params for checkpoints, summaries, and logging
      log_interval=1000,
      summary_interval=1000,
      summaries_flush_secs=10,
      debug_summaries=True,
      summarize_grads_and_vars=True,
      eval_metrics_callback=None,
  ):
      """A simple Atari train and eval for DQN.
  
      Args:
        root_dir: Directory to write log files to.
        env_name: Fully-qualified name of the Atari environment (i.e. Pong-v0).
        num_iterations: Number of train/eval iterations to run.
        max_episode_frames: Maximum length of a single episode, in ALE frames.
        terminal_on_life_loss: Whether to simulate an episode termination when a
          life is lost.
        conv_layer_params: Params for convolutional layers of QNetwork.
        fc_layer_params: Params for fully connected layers of QNetwork.
        initial_collect_steps: Number of frames to ALE frames to process before
          beginning to train. Since this is in ALE frames, there will be
          initial_collect_steps/4 items in the replay buffer when training starts.
        epsilon_greedy: Final epsilon value to decay to for training.
        epsilon_decay_period: Period over which to decay epsilon, from 1.0 to
          epsilon_greedy (defined above).
        replay_buffer_capacity: Maximum number of items to store in the replay
          buffer.
        train_steps_per_iteration: Number of ALE frames to run through for each
          iteration of training.
        update_period: Run a train operation every update_period ALE frames.
        target_update_tau: Coeffecient for soft target network updates (1.0 ==
          hard updates).
        target_update_period: Period, in ALE frames, to copy the live network to
          the target network.
        batch_size: Number of frames to include in each training batch.
        learning_rate: RMS optimizer learning rate.
        n_step_update: The number of steps to consider when computing TD error and
          TD loss. Applies standard single-step updates when set to 1.
        gamma: Discount for future rewards.
        reward_scale_factor: Scaling factor for rewards.
        gradient_clipping: Norm length to clip gradients.
        do_eval: If True, run an eval every iteration. If False, skip eval.
        eval_steps_per_iteration: Number of ALE frames to run through for each
          iteration of evaluation.
        eval_epsilon_greedy: Epsilon value to use for the evaluation policy (0 ==
          totally greedy policy).
        log_interval: Log stats to the terminal every log_interval training steps.
        summary_interval: Write TF summaries every summary_interval training
          steps.
        summaries_flush_secs: Flush summaries to disk every summaries_flush_secs
          seconds.
        debug_summaries: If True, write additional summaries for debugging (see
          dqn_agent for which summaries are written).
        summarize_grads_and_vars: Include gradients in summaries.
        eval_metrics_callback: A callback function that takes (metric_dict,
          global_step) as parameters. Called after every eval with the results of
          the evaluation.
      """
      self._update_period = update_period / ATARI_FRAME_SKIP
      self._train_steps_per_iteration = (
          train_steps_per_iteration / ATARI_FRAME_SKIP
      )
      self._do_eval = do_eval
      self._eval_steps_per_iteration = (
          eval_steps_per_iteration / ATARI_FRAME_SKIP
      )
      self._eval_epsilon_greedy = eval_epsilon_greedy
      self._initial_collect_steps = (
          initial_collect_steps / ATARI_FRAME_SKIP
      )
      self._summary_interval = summary_interval
      self._num_iterations = num_iterations
      self._log_interval = log_interval
      self._eval_metrics_callback = eval_metrics_callback
  
      with gin.unlock_config():
          gin.bind_parameter(
              (
                  "tf_agents.environments.atari_preprocessing."
                  "AtariPreprocessing.terminal_on_life_loss"
              ),
              terminal_on_life_loss,
          )
  
      root_dir = os.path.expanduser(root_dir)
      train_dir = os.path.join(root_dir, "train")
      eval_dir = os.path.join(root_dir, "eval")
  
      train_summary_writer = tf.compat.v2.summary.create_file_writer(
          train_dir, flush_millis=summaries_flush_secs * 1000
      )
      train_summary_writer.set_as_default()
      self._train_summary_writer = train_summary_writer
  
      self._eval_summary_writer = None
      if self._do_eval:
          self._eval_summary_writer = (
              tf.compat.v2.summary.create_file_writer(
                  eval_dir, flush_millis=summaries_flush_secs * 1000
              )
          )
          self._eval_metrics = [
              py_metrics.AverageReturnMetric(
                  name="PhaseAverageReturn", buffer_size=np.inf
              ),
              py_metrics.AverageEpisodeLengthMetric(
                  name="PhaseAverageEpisodeLength", buffer_size=np.inf
              ),
          ]
  
      self._global_step = tf.compat.v1.train.get_or_create_global_step()
      with tf.compat.v2.summary.record_if(
          lambda: tf.math.equal(
              self._global_step % self._summary_interval, 0
          )
      ):
          self._env = suite_atari.load(
              env_name,
              max_episode_steps=max_episode_frames / ATARI_FRAME_SKIP,
              gym_env_wrappers=suite_atari.DEFAULT_ATARI_GYM_WRAPPERS_WITH_STACKING,
          )
          self._env = batched_py_environment.BatchedPyEnvironment(
              [self._env]
          )
  
          observation_spec = tensor_spec.from_spec(
              self._env.observation_spec()
          )
          time_step_spec = ts.time_step_spec(observation_spec)
          action_spec = tensor_spec.from_spec(self._env.action_spec())
  
          with tf.device("/cpu:0"):
              epsilon = tf.compat.v1.train.polynomial_decay(
                  1.0,
                  self._global_step,
                  epsilon_decay_period
                  / ATARI_FRAME_SKIP
                  / self._update_period,
                  end_learning_rate=epsilon_greedy,
              )
  
          with tf.device("/gpu:0"):
              optimizer = tf.compat.v1.train.RMSPropOptimizer(
                  learning_rate=learning_rate,
                  decay=0.95,
                  momentum=0.0,
                  epsilon=0.00001,
                  centered=True,
              )
              categorical_q_net = AtariCategoricalQNetwork(
                  observation_spec,
                  action_spec,
                  conv_layer_params=conv_layer_params,
                  fc_layer_params=fc_layer_params,
              )
              agent = categorical_dqn_agent.CategoricalDqnAgent(
                  time_step_spec,
                  action_spec,
                  categorical_q_network=categorical_q_net,
                  optimizer=optimizer,
                  epsilon_greedy=epsilon,
                  n_step_update=n_step_update,
                  target_update_tau=target_update_tau,
                  target_update_period=(
                      target_update_period
                      / ATARI_FRAME_SKIP
                      / self._update_period
                  ),
                  gamma=gamma,
                  reward_scale_factor=reward_scale_factor,
                  gradient_clipping=gradient_clipping,
                  debug_summaries=debug_summaries,
                  summarize_grads_and_vars=summarize_grads_and_vars,
                  train_step_counter=self._global_step,
              )
  
              self._collect_policy = py_tf_policy.PyTFPolicy(
                  agent.collect_policy
              )
  
              if self._do_eval:
                  self._eval_policy = py_tf_policy.PyTFPolicy(
                      epsilon_greedy_policy.EpsilonGreedyPolicy(
                          policy=agent.policy,
                          epsilon=self._eval_epsilon_greedy,
                      )
                  )
  
              py_observation_spec = self._env.observation_spec()
              py_time_step_spec = ts.time_step_spec(py_observation_spec)
              py_action_spec = policy_step.PolicyStep(
                  self._env.action_spec()
              )
              data_spec = trajectory.from_transition(
                  py_time_step_spec, py_action_spec, py_time_step_spec
              )
              self._replay_buffer = (
                  py_hashed_replay_buffer.PyHashedReplayBuffer(
                      data_spec=data_spec,
                      capacity=replay_buffer_capacity,
                  )
              )
  
          with tf.device("/cpu:0"):
              ds = self._replay_buffer.as_dataset(
                  sample_batch_size=batch_size,
                  num_steps=n_step_update + 1,
              )
              ds = ds.prefetch(4)
              ds = ds.apply(
                  tf.data.experimental.prefetch_to_device("/gpu:0")
              )
  
          with tf.device("/gpu:0"):
              self._ds_itr = tf.compat.v1.data.make_one_shot_iterator(
                  ds
              )
              experience = self._ds_itr.get_next()
              self._train_op = agent.train(experience)
  
              self._env_steps_metric = py_metrics.EnvironmentSteps()
              self._step_metrics = [
                  py_metrics.NumberOfEpisodes(),
                  self._env_steps_metric,
              ]
              self._train_metrics = self._step_metrics + [
                  py_metrics.AverageReturnMetric(buffer_size=10),
                  py_metrics.AverageEpisodeLengthMetric(buffer_size=10),
              ]
              # The _train_phase_metrics average over an entire train iteration,
              # rather than the rolling average of the last 10 episodes.
              self._train_phase_metrics = [
                  py_metrics.AverageReturnMetric(
                      name="PhaseAverageReturn", buffer_size=np.inf
                  ),
                  py_metrics.AverageEpisodeLengthMetric(
                      name="PhaseAverageEpisodeLength",
                      buffer_size=np.inf,
                  ),
              ]
              self._iteration_metric = py_metrics.CounterMetric(
                  name="Iteration"
              )
  
              # Summaries written from python should run every time they are
              # generated.
              with tf.compat.v2.summary.record_if(True):
                  self._steps_per_second_ph = tf.compat.v1.placeholder(
                      tf.float32, shape=(), name="steps_per_sec_ph"
                  )
                  self._steps_per_second_summary = (
                      tf.compat.v2.summary.scalar(
                          name="global_steps_per_sec",
                          data=self._steps_per_second_ph,
                          step=self._global_step,
                      )
                  )
  
                  for metric in self._train_metrics:
                      metric.tf_summaries(
                          train_step=self._global_step,
                          step_metrics=self._step_metrics,
                      )
  
                  for metric in self._train_phase_metrics:
                      metric.tf_summaries(
                          train_step=self._global_step,
                          step_metrics=(self._iteration_metric,),
                      )
                  self._iteration_metric.tf_summaries(
                      train_step=self._global_step
                  )
  
                  if self._do_eval:
                      with self._eval_summary_writer.as_default():
                          for metric in self._eval_metrics:
                              metric.tf_summaries(
                                  train_step=self._global_step,
                                  step_metrics=(
                                      self._iteration_metric,
                                  ),
                              )
  
              self._train_checkpointer = common.Checkpointer(
                  ckpt_dir=train_dir,
                  agent=agent,
                  global_step=self._global_step,
                  optimizer=optimizer,
                  metrics=metric_utils.MetricsGroup(
                      self._train_metrics
                      + self._train_phase_metrics
                      + [self._iteration_metric],
                      "train_metrics",
                  ),
              )
              self._policy_checkpointer = common.Checkpointer(
                  ckpt_dir=os.path.join(train_dir, "policy"),
                  policy=agent.policy,
                  global_step=self._global_step,
              )
              self._rb_checkpointer = common.Checkpointer(
                  ckpt_dir=os.path.join(train_dir, "replay_buffer"),
                  max_to_keep=1,
                  replay_buffer=self._replay_buffer,
              )
  
              self._init_agent_op = agent.initialize()
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  import numpy
  import tensorflow
  from tf_agents.agents.categorical_dqn import categorical_dqn_agent
  from tf_agents.environments import batched_py_environment
  from tf_agents.environments import suite_atari
  from tf_agents.eval import metric_utils
  from tf_agents.metrics import py_metric
  from tf_agents.metrics import py_metrics
  from tf_agents.networks import categorical_q_network
  from tf_agents.networks import network
  from tf_agents.policies import epsilon_greedy_policy
  from tf_agents.policies import py_tf_policy
  from tf_agents.policies import random_py_policy
  from tf_agents.replay_buffers import py_hashed_replay_buffer
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import common
  from tf_agents.utils import timer
label_desc: Use image node_id 1 to create a new TrainEval object from inherited base classes: object with example: obj = TrainEval(root_dir, env_name, num_iterations, max_episode_frames, terminal_on_life_loss, conv_layer_params, fc_layer_params, initial_collect_steps, epsilon_greedy, epsilon_decay_period, replay_buffer_capacity, train_steps_per_iteration, update_period, target_update_tau, target_update_period, batch_size, learning_rate, n_step_update, gamma, reward_scale_factor, gradient_clipping, do_eval, eval_steps_per_iteration, eval_epsilon_greedy, log_interval, summary_interval, summaries_flush_secs, debug_summaries, summarize_grads_and_vars, eval_metrics_callback)