text: import os
  import time
  import tensorflow as tf
  from tf_agents.utils import common
  
  from env import CartPole
  from agent.dqn import DQN
  from buffer import ReplayBuffer
  from eval import ExpectedReturn
  
  # Compulsory config for tf_agents
  tf.compat.v1.enable_v2_behavior()
  
  # Trick
  # No GPU: my super-extra-fast-and-furiuos-ahuhu machine
  # GPUs: training servers
  LOCAL = not len(tf.config.list_physical_devices('GPU')) > 0
  
  # Saving dir
  POLICY_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            './models/policy')
  CHECKPOINT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                './models/checkpoints')
  
  # Environment
  train_env = CartPole.env()
  eval_env = CartPole.env()
  
  # Agent
  dqn = DQN(train_env, CHECKPOINT_DIR)
  dqn.agent.train = common.function(dqn.agent.train)
  
  # Metrics and Evaluation
  criterion = ExpectedReturn()
  
  # Replay buffer
  initial_collect_steps = 100
  replay_buffer = ReplayBuffer(
      dqn.agent.collect_data_spec,
      batch_size=train_env.batch_size,
  )
  for _ in range(initial_collect_steps):
      replay_buffer.collect_steps(train_env, dqn.agent.collect_policy)
  dataset = replay_buffer.as_dataset()
  iterator = iter(dataset)
  
  # Train
  num_iterations = 10000000
  eval_step = 1000
  start = time.time()
  loss = 0
  step = dqn.agent.train_step_counter.numpy()
  while step <= num_iterations:
      if LOCAL:
          train_env.render()
      replay_buffer.collect_steps(train_env, dqn.agent.collect_policy)
      experience, _ = next(iterator)
      loss += dqn.agent.train(experience).loss
      # Evaluation
      step = dqn.agent.train_step_counter.numpy()
      if step % eval_step == 0:
          # Checkpoints
          dqn.save_checkpoint()
          # Evaluation
          avg_return = criterion.eval(eval_env, dqn.agent.policy)
          print('Step = {0}: Average Return = {1} / Average Loss = {2}'.format(
              step, avg_return, loss/eval_step))
          end = time.time()
          print('Step estimated time: {:.4f}'.format((end-start)/eval_step))
          # Reset
          start = time.time()
          loss = 0
  
  # Visualization
  criterion.save()
  