code: def main():
      env = gym.make(args.env)
      env = wrap_deepmind(
          env, dim=IMAGE_SIZE[0], framestack=False, obs_format="NCHW"
      )
      test_env = gym.make(args.env)
      test_env = wrap_deepmind(
          test_env, dim=IMAGE_SIZE[0], obs_format="NCHW", test=True
      )
  
      # set seed
      torch.manual_seed(args.train_seed)
      torch.cuda.manual_seed_all(args.train_seed)
      env.seed(args.train_seed)
      test_env.seed(args.test_seed)
  
      rpm = ReplayMemory(MEMORY_SIZE, IMAGE_SIZE, CONTEXT_LEN)
      act_dim = env.action_space.n
  
      # get model
      model = AtariModel(act_dim, args.dueling)
  
      # get algorithm
      algo_name = args.algo
      if algo_name == "DQN":
          alg = DQN(model, gamma=GAMMA, lr=LR_START)
      else:
          alg = DDQN(model, gamma=GAMMA, lr=LR_START)
  
      # get agent
      agent = AtariAgent(
          alg,
          act_dim=act_dim,
          total_step=TOTAL_STEP,
          update_target_step=UPDATE_TARGET_STEP,
          start_lr=LR_START,
      )
  
      # start training, memory warm up
      with tqdm(
          total=MEMORY_WARMUP_SIZE, desc="[Replay Memory Warm Up]"
      ) as pbar:
          while rpm.size() < MEMORY_WARMUP_SIZE:
              total_reward, steps, _ = run_train_episode(
                  agent, env, rpm
              )
              pbar.update(steps)
  
      test_flag = 0
      train_total_steps = args.train_total_steps
      pbar = tqdm(total=train_total_steps)
      cum_steps = 0  # this is the current timestep
      while cum_steps < train_total_steps:
          # start epoch
          total_reward, steps, loss = run_train_episode(agent, env, rpm)
          cum_steps += steps
  
          pbar.set_description(
              "[train]exploration, learning rate:{}, {}".format(
                  agent.curr_ep,
                  agent.alg.optimizer.param_groups[0]["lr"],
              )
          )
          summary.add_scalar(
              "{}/training_rewards".format(algo_name),
              total_reward,
              cum_steps,
          )
          summary.add_scalar(
              "{}/loss".format(algo_name), loss, cum_steps
          )  # mean of total loss
          summary.add_scalar(
              "{}/exploration".format(algo_name),
              agent.curr_ep,
              cum_steps,
          )
  
          pbar.update(steps)
  
          # perform evaluation
          if cum_steps // args.test_every_steps >= test_flag:
              while cum_steps // args.test_every_steps >= test_flag:
                  test_flag += 1
  
              pbar.write("testing")
              eval_rewards_mean = run_evaluate_episodes(agent, test_env)
  
              logger.info(
                  "eval_agent done, (steps, eval_reward): ({}, {})".format(
                      cum_steps, eval_rewards_mean
                  )
              )
  
              summary.add_scalar(
                  "{}/mean_validation_rewards".format(algo_name),
                  eval_rewards_mean,
                  cum_steps,
              )
  
      pbar.close()
  
file_path: PARL/benchmark/torch/dqn/train.py
imports: import argparse
  import os
  import torch
  import parl
  import gym
  import numpy
  from tqdm import tqdm
  from parl.utils import summary, logger
  from parl.algorithms import DQN, DDQN
  from parl.env.atari_wrappers import wrap_deepmind
  from agent import AtariAgent
  from model import AtariModel
  from replay_memory import ReplayMemory, Experience
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types