text: #!/usr/bin/env python3
  import gym
  import ptan
  import ptan.ignite as ptan_ignite
  from datetime import datetime, timedelta
  import argparse
  import random
  import warnings
  
  import torch
  import torch.optim as optim
  import torch.nn as nn
  
  from ignite.engine import Engine
  from ignite.metrics import RunningAverage
  from ignite.contrib.handlers import tensorboard_logger as tb_logger
  
  from lib import dqn_model, common
  
  NAME = "00_slow_grads"
  
  
  def batch_generator(buffer: ptan.experience.ExperienceReplayBuffer,
                      initial: int, batch_size: int):
      buffer.populate(initial)
      while True:
          buffer.populate(1)
          yield buffer.sample(batch_size)
  
  
  class DQNAgent(ptan.agent.BaseAgent):
      """
      DQNAgent is a memoryless DQN agent which calculates Q values
      from the observations and  converts them into the actions using action_selector
      """
      def __init__(self, dqn_model, action_selector, device="cpu", preprocessor=ptan.agent.default_states_preprocessor):
          self.dqn_model = dqn_model
          self.action_selector = action_selector
          self.preprocessor = preprocessor
          self.device = device
  
      def __call__(self, states, agent_states=None):
          if agent_states is None:
              agent_states = [None] * len(states)
          if self.preprocessor is not None:
              states = self.preprocessor(states)
              if torch.is_tensor(states):
                  states = states.to(self.device)
          q_v = self.dqn_model(states)
          q = q_v.data.cpu().numpy()
          actions = self.action_selector(q)
          return actions, agent_states
  
  
  def calc_loss_dqn(batch, net, tgt_net, gamma, device="cpu", cuda_async=False):
      states, actions, rewards, dones, next_states = common.unpack_batch(batch)
  
      states_v = torch.tensor(states).to(device, non_blocking=cuda_async)
      next_states_v = torch.tensor(next_states).to(device, non_blocking=cuda_async)
      actions_v = torch.tensor(actions).to(device, non_blocking=cuda_async)
      rewards_v = torch.tensor(rewards).to(device, non_blocking=cuda_async)
      done_mask = torch.BoolTensor(dones).to(device, non_blocking=cuda_async)
  
      state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)
      next_state_values = tgt_net(next_states_v).max(1)[0]
      next_state_values[done_mask] = 0.0
  
      expected_state_action_values = next_state_values.detach() * gamma + rewards_v
      return nn.MSELoss()(state_action_values, expected_state_action_values)
  
  
  if __name__ == "__main__":
      # get rid of missing metrics warning
      warnings.simplefilter("ignore", category=UserWarning)
  
      random.seed(common.SEED)
      torch.manual_seed(common.SEED)
      params = common.HYPERPARAMS['pong']
      parser = argparse.ArgumentParser()
      parser.add_argument("--cuda", default=False, action="store_true", help="Enable cuda")
      args = parser.parse_args()
      device = torch.device("cuda" if args.cuda else "cpu")
  
      env = gym.make(params.env_name)
      env = ptan.common.wrappers.wrap_dqn(env)
      env.seed(common.SEED)
  
      net = dqn_model.DQN(env.observation_space.shape, env.action_space.n).to(device)
  
      tgt_net = ptan.agent.TargetNet(net)
      selector = ptan.actions.EpsilonGreedyDQNActionSelector(epsilon=params.epsilon_start)
      epsilon_tracker = common.EpsilonTracker(selector, params)
      agent = DQNAgent(net, selector, device=device)
  
      experience_source = ptan.experience.ExperienceSourceFirstLast(
          env, agent, gamma=params.gamma, steps_count=1)
      buffer = ptan.experience.ExperienceReplayBuffer(
          experience_source, buffer_size=params.replay_size)
      optimizer = optim.Adam(net.parameters(), lr=params.learning_rate)
  
      def process_batch(engine, batch):
          optimizer.zero_grad()
          loss_v = calc_loss_dqn(batch, net, tgt_net.target_model,
                                 gamma=params.gamma, device=device)
          loss_v.backward()
          optimizer.step()
          epsilon_tracker.frame(engine.state.iteration)
          if engine.state.iteration % params.target_net_sync == 0:
              tgt_net.sync()
          return {
              "loss": loss_v.item(),
              "epsilon": selector.epsilon,
          }
  
      engine = Engine(process_batch)
      ptan_ignite.EndOfEpisodeHandler(experience_source, bound_avg_reward=17.0).attach(engine)
      ptan_ignite.EpisodeFPSHandler().attach(engine)
  
      @engine.on(ptan_ignite.EpisodeEvents.EPISODE_COMPLETED)
      def episode_completed(trainer: Engine):
          print("Episode %d: reward=%s, steps=%s, speed=%.3f frames/s, elapsed=%s" % (
              trainer.state.episode, trainer.state.episode_reward,
              trainer.state.episode_steps, trainer.state.metrics.get('fps', 0),
              timedelta(seconds=trainer.state.metrics.get('time_passed', 0))))
  
      @engine.on(ptan_ignite.EpisodeEvents.BOUND_REWARD_REACHED)
      def game_solved(trainer: Engine):
          print("Game solved in %s, after %d episodes and %d iterations!" % (
              timedelta(seconds=trainer.state.metrics['time_passed']),
              trainer.state.episode, trainer.state.iteration))
          trainer.should_terminate = True
  
      logdir = f"runs/{datetime.now().isoformat(timespec='minutes')}-{params.run_name}-{NAME}"
      tb = tb_logger.TensorboardLogger(log_dir=logdir)
      RunningAverage(output_transform=lambda v: v['loss']).attach(engine, "avg_loss")
  
      episode_handler = tb_logger.OutputHandler(tag="episodes", metric_names=['reward', 'steps', 'avg_reward'])
      tb.attach(engine, log_handler=episode_handler, event_name=ptan_ignite.EpisodeEvents.EPISODE_COMPLETED)
  
      # write to tensorboard every 100 iterations
      ptan_ignite.PeriodicEvents().attach(engine)
      handler = tb_logger.OutputHandler(tag="train", metric_names=['avg_loss', 'avg_fps'],
                                        output_transform=lambda a: a)
      tb.attach(engine, log_handler=handler, event_name=ptan_ignite.PeriodEvents.ITERS_100_COMPLETED)
  
      engine.run(batch_generator(buffer, params.replay_initial, params.batch_size))
  