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.')
Run the trainer for n_epochs and call close on it.
trax/rl_trainer.py