code: def train_world_model(
      env, data_dir, output_dir, hparams, world_model_steps_num, epoch
  ):
      """Train the world model on problem_name."""
      world_model_steps_num += world_model_step_increment(
          hparams, epoch
      )
      model_hparams = trainer_lib.create_hparams(
          hparams.generative_model_params
      )
      model_hparams.learning_rate = model_hparams.learning_rate_constant
      if epoch > 0:
          model_hparams.learning_rate *= hparams.learning_rate_bump
      if hparams.wm_policy_param_sharing:
          model_hparams.optimizer_zero_grads = True
  
      restarter = Restarter(
          "world_model", output_dir, world_model_steps_num
      )
      if restarter.should_skip:
          return world_model_steps_num
      with restarter.training_loop():
          train_supervised(
              problem=env,
              model_name=hparams.generative_model,
              hparams=model_hparams,
              data_dir=data_dir,
              output_dir=output_dir,
              train_steps=restarter.target_global_step,
              eval_steps=100,
              local_eval_frequency=2000,
          )
  
      return world_model_steps_num
  
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 10 for calling a global function with example usage: train_world_model(env, data_dir, output_dir, hparams, world_model_steps_num, epoch) and returns: world_model_steps_num, world_model_steps_num