I know the parameters chosen for training a RL model depend heavily on the model itself as well as the problem. Nevertheless, I am trying to train a bunch of these agents on different environments, and would like to do a preliminar training in order to check which Environment-agent combination seems more fitting or promising for my problem to then develop this one more precisely. The thing is that I do not count with advanced resources (I have a dedicated CPU, but no GPU) and the executions take very long, so I must select these parameters very wisely in order to run a training process significant enough but also fast enough to have the results at least in two weeks or three. I'd like to train 3 agents on 4 different environments, and with my current code and parameters, a leaning episode takes 2 hours, while an evaluation episode (and I'm tying to run 5 of them per each learning one) takes about 8 hours, so I think my code must have some kind of flaw that causes it to last so much on the evaluation part...
As you will see in the code, I constantly store model versions. This is because I am using a selenium webdriver script in my agent's step method, and for some reason every now and then (like every hour or so) Chrome DevTools disconnect, to which the only solution I have found has been to store constant copies of the model and the rewards in order to be able to continue from the last checkpoint every time I rerun the training (to be able to not start from the beginning over and over again)
Therefore, I'd like to know:
Which would be a good balance of episodes-iterations for learning and evaluation? I am currently doing this:
- The environment returns "done" when a certain condition is fulfilled or after a maximum of 1024 iterations
- Training episodes: 10
- Learning iterations per training episode: 1024
- Evaluation episodes per training step: 5
- Evaluation iterations per evaluation episode: 1024
- Iteration steps per to store evaluation checkpoints: 103
Is there anything I can change in my current code to make all these executions much more efficient? Because I think I must be doing something wrong for it to last for so long (specially on evaluation)
Any other suggestion to improve the training process or the code?
def learn(self, total_timesteps, load_if_previous = False, training_step = 0, force_more_trainig = False, checkpoint_interval=100, callback=None, log_interval=1, tb_log_name='AgentV1',
reset_num_timesteps=True, progress_bar=True):
print("\n\t---------------------------------------------------------------------------------------")
print(f"\t\t\t\t\t\t LEARNING {datetime.datetime.now()}")
print("\t---------------------------------------------------------------------------------------")
current_step = 0
final_model_name = f"{self.model_path}_{training_step}_final"
# Check if the model already exists
if os.path.exists(f"{final_model_name}.zip") and load_if_previous and not force_more_trainig:
# Skip learning step o
print(f"\tModel for training step {training_step} already exists ({final_model_name}.zip). Loading the model.")
else:
while current_step < total_timesteps:
remaining_steps = total_timesteps - current_step
steps_to_train = min(checkpoint_interval, remaining_steps)
print(f"\t ------ Training from {current_step} to {current_step + steps_to_train} ------ ")
# Check if the model already exists
iternmediate_model_file = f"{self.model_path}_{training_step}_{current_step}.zip"
print(f"\t\t Finding {iternmediate_model_file}...")
if os.path.exists(iternmediate_model_file) and load_if_previous:
print(f"\t\t Model for training step and batch {training_step} already exists ({iternmediate_model_file}.zip). Loading the model")
self.load(iternmediate_model_file)
else:
print(f"\t\t Learning batch #{current_step} of {total_timesteps}")
# Set the model's environment before training
self.model.set_env(self.gym_env )
self.model.learn(total_timesteps=steps_to_train, callback=callback, log_interval=log_interval,
tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps, progress_bar=progress_bar)
# Save the model after every checkpoint interval
if current_step % checkpoint_interval == 0:
self.save(f"{self.model_path}_{training_step}_{current_step}")
current_step += steps_to_train
# Remove non-final intermediate files
for intermediate_file in glob.glob(f"{self.model_path}_{training_step}_*.zip"):
if intermediate_file != f"{final_model_name}.zip" and intermediate_file != f"{final_model_name}":
os.remove(intermediate_file)
print(f"\t\t - Deleted intermediate file: {intermediate_file}")
# Save the final model
self.save(final_model_name)
print(f"\t--------------------- ENDOF Learning {datetime.datetime.now()} ---------------------\n")
return self.load(f"{final_model_name}.zip")
def evaluate(self, checkpoint_interval = 100, deterministic=False, training_step = 0):
"""
The model will be evaluated for self.n_eval_episodes times until "is done"
"""
n_batch_iterations = math.ceil(self.gym_env.max_steps / checkpoint_interval)
print("\n\t---------------------------------------------------------------------------------------")
print(f"\t\t\t\t\t\t EVALUATION {datetime.datetime.now()}")
print("\t---------------------------------------------------------------------------------------")
print(f"\t\t - Evaluation episodes: {self.n_eval_episodes}")
print(f"\t\t - Iterations for the step: {self.gym_env.max_steps}")
print(f"\t\t - Checkpoints for the step: {checkpoint_interval}")
# ------------------------------------- Checkpoint Loadng -------------------------------------
# LOAD last execution's episode, batch and agent information
checkpoint_files = glob.glob(f"{self.model_path}_evaluation_Teps_{training_step}_episode_*_batch_*.zip")
checkpoint_files.sort()
if len(checkpoint_files) > 0:
last_checkpoint_file = checkpoint_files[-1]
episode_str, batch_str = last_checkpoint_file.split(f"_evaluation_Teps_{training_step}_episode_")[1].split("_batch_")
last_episode = int(episode_str)
# Obtain the batch from the string (if errors, batch = 0)
error = True
attempts = 5
loops = 0
last_batch = 0
while error or loops < attempts:
loops += 1
try:
batch_str = batch_str[:-4]
last_batch = int(batch_str)
error = False
except ValueError as e:
pass
print(f"\t\t Found checkpoint file: {last_checkpoint_file}")
print(f"\t\t - Last episode: {last_episode}")
print(f"\t\t - Last batch: {last_batch}")
# Load the file's information into "self"
self.load(last_checkpoint_file)
else:
print("\t\t - No checkpoint file found. Starting evaluation from scratch.")
last_episode = 0
last_batch = 0
# Create a temporary environment to evaluate the policy
eval_env = self.gym_env
reward_file_path = f"{self.monitor_gym_log_folder}/Rewards_{self.model_name}_Evaluation_Teps_{self.gym_env.current_episode}.csv"
# Initialize lists to store rewards
episode_rewards = []
# LOAD episode rewards from previous executions
episode_rewards_file = f"{self.model_path}_evaluation_rewards_for_training_step_{training_step}.npy"
if os.path.exists(episode_rewards_file):
print(f"\t\t Rewards loaded from: {episode_rewards_file}")
episode_rewards = np.load(episode_rewards_file)
# ------------------------------------- ENDOF Loading -------------------------------------
batch = last_batch
iteration_count = 0
step = 0
for episode in range(last_episode, self.n_eval_episodes):
print(f"\n\t\t Executing evaluation EPISODE #{episode + 1}/{self.n_eval_episodes + 1}")
obs = eval_env.reset()[0]
done = False
episode_reward = episode_rewards[episode] if episode < len(episode_rewards) else 0.0
print(f"\n\t\t Executing evaluation BATCH #{batch}/{n_batch_iterations}")
while not done:
step += 1
# Get the action from the model
action, _ = self.model.predict(obs, deterministic=deterministic)
# Perform the action in the environment
obs, reward, done, truncated, info = eval_env.step(action, episode = episode + 1, step = step, reward_file_path = reward_file_path)
# obs, reward, done, info = eval_env.step(action, episode = episode, step = iteration_count, reward_file_path = reward_file_path)
# Accumulate the reward
episode_reward = reward
# Increment the iteration count in order to save (or not) a new checkpoint
iteration_count += 1
if iteration_count > checkpoint_interval:
# ------------------------------------- Checkpoint Storage -------------------------------------
# After each batch, store the information about batch, episode and the agent
print(f"\t\t STORING BATCH: info Batch checkpoint stored at {self.model_path}_evaluation_Teps_{training_step}_episode_{episode}_batch_{batch}.zip")
self.save(f"{self.model_path}_evaluation_Teps_{training_step}_episode_{episode}_batch_{batch}.zip")
# Save the rewards after that batch
print(f"\t\t STORING REWARDS")
np.save(episode_rewards_file, episode_rewards)
# ------------------------------------- ENDOF Checkpoint Storage -------------------------------------
batch += 1
iteration_count = 0
print(f"\n\t\t Executing evaluation BATCH #{batch}/{n_batch_iterations}")
done = False
step = 0
if episode < len(episode_rewards):
episode_rewards[episode] = episode_reward
else:
if isinstance(episode_rewards, np.ndarray):
episode_rewards = np.append(episode_rewards, episode_reward)
else:
episode_rewards.append(episode_reward)
episode_reward = 0
# ------------------------------------------------------------ Checkpoint Deletion ------------------------------------------------------------
# After each episode, delete the intermidiate file where episode is lower than current episode (the same for rewards)
# Delete intermediate files for episodes and batches lower than the current episode
for episode_file in glob.glob(f"{self.model_path}_evaluation_Teps_{training_step}_episode_*.zip"):
episode_num = int(episode_file.split(f"_evaluation_Teps_{training_step}_episode_")[1].split("_batch")[0])
if episode_num < episode:
os.remove(episode_file)
print(f"\t\t - Deleted intermediate episode file: {episode_file}")
for batch_file in glob.glob(f"{self.model_path}_evaluation_Teps_{training_step}_episode_{episode}_batch_*.zip"):
batch_num = int(batch_file.split("_batch_")[1][:-4])
if batch_num < batch:
os.remove(batch_file)
print(f"\t\t - Deleted intermediate batch file: {batch_file}")
# ------------------------------------------------------------ ENDOF Checkpoint Deletion ------------------------------------------------------
if episode >= self.n_eval_episodes:
break
# Calculate mean and standard deviation of rewards
mean_reward = np.mean(episode_rewards)
std_reward = np.std(episode_rewards)
print(f"\t\t Mean reward = {mean_reward:.2f} +/- {std_reward}")
print(f"\t--------------------- ENDOF Evaluation {datetime.datetime.now()} ---------------------\n")
return mean_reward
def choose_action(self, obs, deterministic = False):
return self.model.predict(obs, deterministic=deterministic)
def train_agent(self, n_steps = 2048, total_timesteps = 50000, callback=None, log_interval=1, tb_log_name='AgentV1', reset_num_timesteps=True, progress_bar=False, checkpoint = True, force_more_trainig = False):
# Set the model's environment before training
self.model.set_env(self.gym_env )
base_path = f"{self.model_path}_{total_timesteps}_{n_steps}_"
extension = ".zip"
plot_files = []
print("--------------------------------------------------------------------------------------")
print(f" Training agent {self.model_name}")
print("--------------------------------------------------------------------------------------\n")
print(f" Training info:", flush=True)
print(f" - Total training steps: {n_steps}", flush=True)
print(f" - Iterations per step: {total_timesteps}", flush=True)
print(f" - Model stored as: {base_path}_*{extension}")
mean_reward = 0
if self.trained and not force_more_trainig:
print("A trained model has already been found, training will be skipped.")
print(f"If you want to keep training the trained model {self.model_name}, rerun the training function by adding force_more_trainig = True")
else:
# train the agent
# Loading the last training info to continue from the same point in case the execution ends abruptly
if checkpoint:
best_mean_reward, best_model_path = -float('inf'), None
model_paths = glob.glob(f"{base_path}_*{extension}")
# Check if there is a saved model
if model_paths:
best_model_path = max(model_paths, key=lambda path: int(path.split('_')[-1].split('.')[0]))
print("Loading saved model...")
self.load(best_model_path)
last_step = int(best_model_path.split("_")[-1].split('.')[0])
print(f"Continuing training from step #{last_step}")
else:
last_step = 0
best_mean_reward, best_model_path = -float('inf'), None
for i in range(1, n_steps+1):
print(f"\n_______________________________________________________________Training EPISODE #{i}/{n_steps}_______________________________________________________________", flush=True)
self.gym_env.reset()
print(f"\tIterations for the step: {total_timesteps}", flush=True)
self.gym_env.set_test_case(data_wrapper.insert("Test_cases", {
"Execution_uid": self.gym_env.execution['uid'],
"Initial_time": utils.getDateTime(),
"Error_count": 0,
"Unique_error_count": 0,
}))
initial_time = utils.getDateTime()
self.learn(total_timesteps=total_timesteps, load_if_previous = True, force_more_trainig = force_more_trainig, training_step = i)
# evaluate the agent and log the results
# If a model for the next step already exists, skip evaluation of the current one
if len(glob.glob(f"{self.model_path}_{i + 1}_*.zip")) > 0:
print(f"\t\t Evaluation for the current Episode has already been completed")
else:
mean_reward = self.evaluate(checkpoint_interval = math.ceil(self.gym_env.max_steps / 10) , deterministic=False, training_step = i)
if mean_reward > best_mean_reward:
best_mean_reward = mean_reward
best_model_path = f"{base_path}_{i}{extension}"
self.save(best_model_path)
print(f"\tNew best model saved to {best_model_path}")
print(f"\tNew best model saved to {best_model_path}", flush=True)
data_wrapper.update("Test_cases", {"uid": self.gym_env.test_case['uid'] }, { "$set": self.gym_env.test_case })
data_wrapper.update("Test_cases", {"uid": self.gym_env.test_case['uid'] }, { "$set": {"End_time": utils.getDateTime(), 'Execution_uid': self.gym_env.execution['uid'], 'Initial_time': initial_time }})
# Generate plots after each episode
if self.monitor:
try:
print("Generating plot to send...")
# Call plot_results and redirect the output to a StringIO object
results_plotter.plot_results([ self.monitor_gym_log_folder ], total_timesteps + 1, results_plotter.X_TIMESTEPS, f"{self.model_name} training for episode {i}")
# Rename the monitor file in self.monitor_gym_log_folder/monitor.csv to self.monitor_gym_log_folder/monitor_{i}.csv by making a copy
try:
monitor_file_path = os.path.join(self.monitor_gym_log_folder, "monitor.csv")
renamed_monitor_file_path = os.path.join(self.monitor_gym_log_folder, f"monitor_{i}.csv")
shutil.copy(monitor_file_path, renamed_monitor_file_path)
except Exception as e:
print("An error occurred when trying to rename the monitor file")
print(e)
# Capture the plot
fig = plt.gcf()
# Save the figure to a file
plot_file = f"{self.monitor_gym_log_folder}/training_results_{self.model_name}_episode_{i}_{ str(int(time.time())) }.png"
# Save the plot to a file
fig.savefig(plot_file)
print(f"Plot file saved at {plot_file}")
plot_files.append(plot_file)
plt.close(fig) # Close the plot to free up memory
except Exception as e:
print("Error when generating the plot with the built in function. ERROR:")
print(e)
print("Plot information: ")
print(f" - Source Monitor file name: {self.monitor_gym_log_folder}/monitor.csv")
print(f" - Renamed Monitor file name: {self.monitor_gym_log_folder}/monitor_{i}.csv")
print(f" - Target file name: {self.monitor_gym_log_folder}/training_results_{self.model_name}_episode_{i}_{ str(int(time.time())) }.png")
traceback.print_exc()
plot_file = ''
print(f"\n_________________________________________ Training step #{i}/{n_steps} ended at: {datetime.datetime.now()} _________________________________________\n\n\n", flush=True)
# SAFE THE FINAL MODEL
self.save(self.model_path)
# Modify the execution to save the model name and mean reward
self.gym_env.execution["Mean_reward"] = mean_reward
self.gym_env.execution["End_time"] = utils.getDateTime()
data_wrapper.update("Executions", {"uid": self.gym_env.execution['uid'] }, { "$set": self.gym_env.execution })
# Send email and plot
utils.sendTrainingEndEmail( self.gym_env.execution['uid'], tb_log_name, mean_reward, plot_files, self.gym_env.test_case)
```
