code: def main():
      # set training env and test env
      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
      )
  
      env.seed(args.train_seed)
      test_env.seed(args.test_seed)
  
      act_dim = env.action_space.n
      algo_name = args.algo
      rpm = ReplayMemory(MEMORY_SIZE, IMAGE_SIZE, CONTEXT_LEN)
  
      # build model
      model = AtariModel(act_dim=act_dim, dueling=args.dueling)
  
      # get algorithm
      if algo_name == "DQN":
          alg = DQN(model, gamma=GAMMA, lr=LR_START)
      else:
          alg = DDQN(model, gamma=GAMMA, lr=LR_START)
  
      # build Agent using model and algorithm
      agent = AtariAgent(
          alg,
          act_dim,
          LR_START,
          args.train_total_steps // 10,
          UPDATE_TARGET_STEP,
      )
  
      # start training, memory warm up
      with tqdm(
          total=args.warmup_size, desc="[Replay Memory Warm Up]"
      ) as pbar:
          while rpm.size() < args.warmup_size:
              total_reward, steps, _ = run_train_episode(
                  agent, env, rpm
              )
              pbar.update(steps)
  
      test_flag = 0
      pbar = tqdm(total=args.train_total_steps)
      cum_steps = 0  # this is the current timestep
      while cum_steps < args.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, alg.optimizer.get_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,
          )
          summary.add_scalar(
              "{}/learning_rate".format(algo_name),
              alg.optimizer.get_lr(),
              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()
  
      # save the parameters to ./model.ckpt
      save_path = "./model/model.ckpt"
      agent.save(save_path)
  
file_path: PARL/examples/DQN_variant/train.py
imports: import os
  import gym
  import numpy
  import paddle
  import parl
  from parl.utils import logger, summary
  from atari_model import AtariModel
  from atari_agent import AtariAgent
  from replay_memory import ReplayMemory, Experience
  from parl.env.atari_wrappers import wrap_deepmind
  from parl.algorithms import DQN, DDQN
  from tqdm import tqdm
  import argparse
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types