def train_rl(
    output_dir,
    n_epochs=10000,
    light_rl=True,
    light_rl_trainer=light_trainers.PolicyGradient):
  """Train the RL agent.

  Args:
    output_dir: Output directory.
    n_epochs: Number epochs to run the training for.
    light_rl: deprecated, always True, left out for old gin configs.
    light_rl_trainer: which light RL trainer to use (experimental).
  """
  del light_rl
  tf_np.set_allow_float64(FLAGS.tf_allow_float64)
  task = rl_task.RLTask()
  env_name = task.env_name


  if FLAGS.jax_debug_nans:
    config.update('jax_debug_nans', True)

  if FLAGS.use_tpu:
    config.update('jax_platform_name', 'tpu')
  else:
    config.update('jax_platform_name', '')


  trainer = light_rl_trainer(task=task, output_dir=output_dir)
  def light_training_loop():
    """Run the trainer for n_epochs and call close on it."""
    try:
      logging.info('Starting RL training for %d epochs.', n_epochs)
      trainer.run(n_epochs, n_epochs_is_total_epochs=True)
      logging.info('Completed RL training for %d epochs.', n_epochs)
      trainer.close()
      logging.info('Trainer is now closed.')
    except Exception as e:
      raise e
    finally:
      logging.info('Encountered an exception, still calling trainer.close()')
      trainer.close()
      logging.info('Trainer is now closed.')

  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
    fastmath.disable_jit()
    with jax.disable_jit():
      light_training_loop()
  else:
    light_training_loop()
Train the RL agent.

  Args:
    output_dir: Output directory.
    n_epochs: Number epochs to run the training for.
    light_rl: deprecated, always True, left out for old gin configs.
    light_rl_trainer: which light RL trainer to use (experimental).
  
trax/rl_trainer.py