code: def train_agent(real_env, learner, world_model_dir, hparams, epoch):
    
    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_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
    )
docs: Train the PPO agent in the simulated environment.