code: def train_world_model(
      env, data_dir, output_dir, hparams, world_model_steps_num, epoch
  ):
    
    world_model_steps_num += world_model_step_increment(
        hparams, is_initial_epoch=(epoch == 0)
    )
    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
docs: Train the world model on problem_name.