code: def train_agent(real_env, learner, world_model_dir, hparams, epoch):
      """Train the PPO agent in the simulated environment."""
      initial_frame_chooser = rl_utils.make_initial_frame_chooser(
          real_env,
          hparams.frame_stack_size,
          hparams.simulation_random_starts,
          hparams.simulation_flip_first_random_for_beginning,
      )
      env_fn = rl.make_simulated_env_fn_from_hparams(
          real_env,
          hparams,
          batch_size=hparams.simulated_batch_size,
          initial_frame_chooser=initial_frame_chooser,
          model_dir=world_model_dir,
          sim_video_dir=os.path.join(
              learner.agent_model_dir, "sim_videos_{}".format(epoch)
          ),
      )
      base_algo_str = hparams.base_algo
      train_hparams = trainer_lib.create_hparams(
          hparams.base_algo_params
      )
      if hparams.wm_policy_param_sharing:
          train_hparams.optimizer_zero_grads = True
  
      rl_utils.update_hparams_from_hparams(
          train_hparams, hparams, base_algo_str + "_"
      )
  
      final_epoch = hparams.epochs - 1
      is_special_epoch = (epoch + 3) == final_epoch or (
          epoch + 7
      ) == final_epoch
      is_special_epoch = is_special_epoch or (
          epoch == 1
      )  # Make 1 special too.
      is_final_epoch = epoch == final_epoch
      env_step_multiplier = (
          3 if is_final_epoch else 2 if is_special_epoch else 1
      )
      learner.train(
          env_fn,
          train_hparams,
          simulated=True,
          save_continuously=True,
          epoch=epoch,
          env_step_multiplier=env_step_multiplier,
      )
  
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 8 for calling a global function with example usage: train_agent(real_env, learner, world_model_dir, hparams, epoch) without return types