keras_rl_patient_training_cb.py

from keras.callbacks import Callback

class PatientTrainingCallback(Callback):
	def __init__(self):
		Callback.__init__(self)
		self.episode_finished = True
		self.end_train_before_episode_end = False

	def on_train_begin(self, logs=None):
		pass

	def on_epoch_begin(self, epoch, logs=None):
		self.episode_finished = False

	def on_batch_begin(self, batch, logs=None):
		pass

	def on_batch_end(self, batch, logs=None):
		pass

	def on_epoch_end(self, epoch, logs=None):
		self.episode_finished = True
		if self.end_train_before_episode_end: self.model.stop_training = True

	def on_train_end(self, logs=None):
		if not self.episode_finished: self.end_train_before_episode_end = True
		self.model.stop_training = False