code: def training_loop(
      hparams, output_dir, report_fn=None, report_metric=None
  ):
      """Run the main training loop."""
      if report_fn:
          assert report_metric is not None
  
      # Directories
      subdirectories = [
          "data",
          "tmp",
          "world_model",
          ("world_model", "debug_videos"),
          "policy",
          "eval_metrics",
      ]
      directories = setup_directories(output_dir, subdirectories)
  
      epoch = -1
      data_dir = directories["data"]
      env = rl_utils.setup_env(
          hparams,
          batch_size=hparams.real_batch_size,
          max_num_noops=hparams.max_num_noops,
          rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
      )
      env.start_new_epoch(epoch, data_dir)
  
      if hparams.wm_policy_param_sharing:
          policy_model_dir = directories["world_model"]
      else:
          policy_model_dir = directories["policy"]
      learner = rl_utils.LEARNERS[hparams.base_algo](
          hparams.frame_stack_size,
          policy_model_dir,
          policy_model_dir,
          hparams.epochs,
      )
  
      # Timing log function
      log_relative_time = make_relative_timing_fn()
  
      # Per-epoch state
      epoch_metrics = []
      metrics = {}
  
      # Collect data from the real environment.
      policy_model_dir = directories["policy"]
      tf.logging.info(
          "Initial training of the policy in real environment."
      )
      train_agent_real_env(env, learner, hparams, epoch)
      metrics[
          "mean_reward/train/clipped"
      ] = rl_utils.compute_mean_reward(
          env.current_epoch_rollouts(), clipped=True
      )
      tf.logging.info(
          "Mean training reward (initial): {}".format(
              metrics["mean_reward/train/clipped"]
          )
      )
      env.generate_data(data_dir)
  
      eval_metrics_writer = tf.summary.FileWriter(
          directories["eval_metrics"]
      )
  
      world_model_steps_num = 0
  
      for epoch in range(hparams.epochs):
          log = make_log_fn(epoch, log_relative_time)
  
          # Train world model
          log("Training world model")
          world_model_steps_num = train_world_model(
              env,
              data_dir,
              directories["world_model"],
              hparams,
              world_model_steps_num,
              epoch,
          )
  
          # Train agent
          log("Training policy in simulated environment.")
          train_agent(
              env, learner, directories["world_model"], hparams, epoch
          )
  
          env.start_new_epoch(epoch, data_dir)
  
          # Train agent on real env (short)
          log("Training policy in real environment.")
          train_agent_real_env(env, learner, hparams, epoch)
  
          if hparams.stop_loop_early:
              return 0.0
  
          env.generate_data(data_dir)
  
          metrics = load_metrics(directories["eval_metrics"], epoch)
          if metrics:
              # Skip eval if metrics have already been written for this epoch. Otherwise
              # we'd overwrite them with wrong data.
              log("Metrics found for this epoch, skipping evaluation.")
          else:
              metrics[
                  "mean_reward/train/clipped"
              ] = rl_utils.compute_mean_reward(
                  env.current_epoch_rollouts(), clipped=True
              )
              log(
                  "Mean training reward: {}".format(
                      metrics["mean_reward/train/clipped"]
                  )
              )
  
              eval_metrics = rl_utils.evaluate_all_configs(
                  hparams, policy_model_dir
              )
              log(
                  "Agent eval metrics:\n{}".format(
                      pprint.pformat(eval_metrics)
                  )
              )
              metrics.update(eval_metrics)
  
              if hparams.eval_world_model:
                  debug_video_path = os.path.join(
                      directories["world_model", "debug_videos"],
                      "{}.avi".format(env.current_epoch),
                  )
                  wm_metrics = rl_utils.evaluate_world_model(
                      env,
                      hparams,
                      directories["world_model"],
                      debug_video_path,
                  )
                  log(
                      "World model eval metrics:\n{}".format(
                          pprint.pformat(wm_metrics)
                      )
                  )
                  metrics.update(wm_metrics)
  
              rl_utils.summarize_metrics(
                  eval_metrics_writer, metrics, epoch
              )
  
              # Report metrics
              if report_fn:
                  if report_metric == "mean_reward":
                      metric_name = rl_utils.get_metric_name(
                          sampling_temp=hparams.eval_sampling_temps[0],
                          max_num_noops=hparams.eval_max_num_noops,
                          clipped=False,
                      )
                      report_fn(eval_metrics[metric_name], epoch)
                  else:
                      report_fn(eval_metrics[report_metric], epoch)
  
          epoch_metrics.append(metrics)
  
      # Return the evaluation metrics from the final epoch
      return epoch_metrics[-1]
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import datetime
  import math
  import os
  import pprint
  import random
  import time
  import six
  from tensor2tensor.bin import t2t_trainer
  from tensor2tensor.models.research import rl
  from tensor2tensor.rl import rl_utils
  from tensor2tensor.rl import trainer_model_based_params
  from tensor2tensor.rl.dopamine_connector import DQNLearner
  from tensor2tensor.rl.restarter import Restarter
  from tensor2tensor.utils import trainer_lib
  import tensorflow.compat.v1
label_desc: Use image node_id 12 for calling a global function with example usage: training_loop(hparams, output_dir, report_fn, report_metric) and returns: epoch_metrics, int