diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 6fa667c..a4e83be 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -130,8 +130,10 @@ class A2C(PPO): logger.logkv("std", th.exp(self.policy.log_std).mean().item()) def learn(self, total_timesteps, callback=None, log_interval=100, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="A2C", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, + tb_log_name="A2C", eval_log_path=None, reset_num_timesteps=True): return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval, eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes, - tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps) + tb_log_name=tb_log_name, eval_log_path=eval_log_path, + reset_num_timesteps=reset_num_timesteps) diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 1e26d15..2e3be4d 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -100,9 +100,10 @@ class CEMRL(TD3): elitism=self.elitism) def learn(self, total_timesteps, callback=None, log_interval=4, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, + tb_log_name="CEMRL", eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) actor_steps = 0 continue_training = True @@ -155,21 +156,6 @@ class CEMRL(TD3): # Get the params back in the population self.es_params[i] = self.actor.parameters_to_vector() - # Evaluate agent - if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: - timesteps_since_eval %= eval_freq - - self.actor.load_from_vector(self.es.mu) - sync_envs_normalization(self.env, eval_env) - - mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes) - evaluations.append(mean_reward) - - if self.verbose > 0: - print("Eval num_timesteps={}, " - "episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward)) - print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time))) - actor_steps = 0 # evaluate all actors for params in self.es_params: @@ -180,7 +166,6 @@ class CEMRL(TD3): n_steps=-1, action_noise=self.action_noise, deterministic=False, callback=callback, learning_starts=self.learning_starts, - num_timesteps=self.num_timesteps, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, log_interval=log_interval) @@ -192,8 +177,6 @@ class CEMRL(TD3): break episode_num += n_episodes - self.num_timesteps += episode_timesteps - timesteps_since_eval += episode_timesteps actor_steps += episode_timesteps self.fitnesses.append(episode_reward) @@ -202,7 +185,6 @@ class CEMRL(TD3): self._update_current_progress(self.num_timesteps, total_timesteps) self.es.tell(self.es_params, self.fitnesses) - timesteps_since_eval += actor_steps callback.on_training_end() diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 5bc2abd..6cfdeb2 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -2,8 +2,7 @@ import time import os import io import zipfile -import typing -from typing import Union, Type, Optional, Dict, Any, List, Tuple +from typing import Union, Type, Optional, Dict, Any, List, Tuple, Callable from abc import ABC, abstractmethod from collections import deque @@ -14,16 +13,13 @@ import numpy as np from torchy_baselines.common import logger from torchy_baselines.common.policies import BasePolicy, get_policy_from_name from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate -from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization +from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize from torchy_baselines.common.monitor import Monitor -from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.save_util import data_to_json, json_to_data from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict +from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback from torchy_baselines.common.noise import ActionNoise -if typing.TYPE_CHECKING: - from torchy_baselines.common.callbacks import BaseCallback - class BaseRLModel(ABC): """ @@ -50,11 +46,12 @@ class BaseRLModel(ABC): :param sde_sample_freq: Sample a new noise matrix every n steps when using SDE Default: -1 (only sample at the beginning of the rollout) """ + def __init__(self, policy: Type[BasePolicy], env: Union[GymEnv, str], policy_base: Type[BasePolicy], - policy_kwargs : Dict[str, Any] = None, + policy_kwargs: Dict[str, Any] = None, verbose: int = 0, device: Union[th.device, str] = 'auto', support_multi_env: bool = False, @@ -133,9 +130,6 @@ class BaseRLModel(ABC): def _get_eval_env(self, eval_env: Optional[GymEnv]) -> Optional[GymEnv]: """ Return the environment that will be used for evaluation. - - :param eval_env: - :return: """ if eval_env is None: eval_env = self.eval_env @@ -146,34 +140,10 @@ class BaseRLModel(ABC): assert eval_env.num_envs == 1 return eval_env - # Type hint as string to avoid circular import - def _init_callback(self, callback) -> 'BaseCallback': - """ - Note: we cannot use type hint here because of circular import. - - :param callback: (Union[callable, [BaseCallback], BaseCallback, None]) - :return: (BaseCallback) - """ - # Avoid circular import - from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback - - # Convert a list of callbacks into a callback - if isinstance(callback, list): - callback = CallbackList(callback) - # Convert functional callback to object - if not isinstance(callback, BaseCallback): - callback = ConvertCallback(callback) - - callback.init_callback(self) - return callback - def scale_action(self, action: np.ndarray) -> np.ndarray: """ Rescale the action from [low, high] to [-1, 1] (no need for symmetric action space) - - :param action: - :return: """ low, high = self.action_space.low, self.action_space.high return 2.0 * ((action - low) / (high - low)) - 1.0 @@ -182,9 +152,6 @@ class BaseRLModel(ABC): """ Rescale the action from [-1, 1] to [low, high] (no need for symmetric action space) - - :param scaled_action: - :return: """ low, high = self.action_space.low, self.action_space.high return low + (0.5 * (scaled_action + 1.0) * (high - low)) @@ -291,7 +258,7 @@ class BaseRLModel(ABC): return self.policy.state_dict() @abstractmethod - def get_opt_parameters(self)-> OptimizerStateDict: + def get_opt_parameters(self) -> OptimizerStateDict: """ Get current model optimizer parameters as dictionary of variable names -> tensors :return: (dict) Dictionary of variable name -> tensor of model's optimizer parameters @@ -300,11 +267,13 @@ class BaseRLModel(ABC): @abstractmethod def learn(self, total_timesteps: int, - callback=None, log_interval: int = 100, + callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, + log_interval: int = 100, tb_log_name: str = "run", eval_env: Optional[GymEnv] = None, eval_freq: int = -1, n_eval_episodes: int = 5, + eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True): """ Return a trained model. @@ -318,6 +287,8 @@ class BaseRLModel(ABC): :param eval_env: (gym.Env) Environment that will be used to evaluate the agent :param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little) :param n_eval_episodes: (int) Number of episode to evaluate the agent + :param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved + :param reset_num_timesteps: (bool) :return: (BaseRLModel) the trained model """ raise NotImplementedError() @@ -391,7 +362,8 @@ class BaseRLModel(ABC): @staticmethod def _load_from_file(load_path: str, load_data: bool = True) -> (Tuple[Optional[Dict[str, Any]], - Optional[TensorDict], Optional[OptimizerStateDict]]): + Optional[TensorDict], + Optional[OptimizerStateDict]]): """ Load model data from a .zip archive :param load_path: Where to load the model from @@ -473,14 +445,53 @@ class BaseRLModel(ABC): if self.eval_env is not None: self.eval_env.seed(seed) - def _setup_learn(self, eval_env: Optional[GymEnv], callback=None) -> (Tuple[int, int, - List[Any], np.ndarray, Optional[VecEnv], Any]): + def _init_callback(self, + callback: Union[None, Callable, List[BaseCallback], BaseCallback], + eval_env: Optional[VecEnv] = None, + eval_freq: int = 10000, + n_eval_episodes: int = 5, + log_path: Optional[str] = None) -> BaseCallback: + """ + :param callback: (Union[callable, [BaseCallback], BaseCallback, None]) + :return: (BaseCallback) + """ + # Convert a list of callbacks into a callback + if isinstance(callback, list): + callback = CallbackList(callback) + + # Convert functional callback to object + if not isinstance(callback, BaseCallback): + callback = ConvertCallback(callback) + + # Create eval callback in charge of the evaluation + if eval_env is not None: + # Same folder as the rest + best_model_save_path = os.path.dirname(log_path) if log_path is not None else None + + eval_callback = EvalCallback(eval_env, + best_model_save_path=best_model_save_path, + log_path=log_path, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes) + callback = CallbackList([callback, eval_callback]) + + callback.init_callback(self) + return callback + + def _setup_learn(self, + eval_env: Optional[GymEnv], + callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, + eval_freq: int = 10000, + n_eval_episodes: int = 5, + log_path: Optional[str] = None + ) -> Tuple[int, np.ndarray, BaseCallback]: """ Initialize different variables needed for training. :param eval_env: (Optional[GymEnv]) :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]]) - :return: (int, int, [float], np.ndarray, VecEnv, BaseCallback) + :param eval_freq: (int) + :param n_eval_episodes: (int) + :param log_path (Optional[str]): + :return: (Tuple[int, np.ndarray, BaseCallback]) """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) @@ -488,18 +499,18 @@ class BaseRLModel(ABC): if self.action_noise is not None: self.action_noise.reset() - callback = self._init_callback(callback) - timesteps_since_eval, episode_num = 0, 0 - evaluations = [] if eval_env is not None and self.seed is not None: eval_env.seed(self.seed) eval_env = self._get_eval_env(eval_env) - obs = self.env.reset() # type: GymEnv + obs = self.env.reset() - return timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback + # Create eval callback if needed + callback = self._init_callback(callback, eval_env, eval_freq, n_eval_episodes, log_path) + + return episode_num, obs, callback def _update_info_buffer(self, infos: List[Dict[str, Any]]) -> None: """ @@ -521,7 +532,6 @@ class BaseRLModel(ABC): action_noise: Optional[ActionNoise] = None, deterministic: bool = False, learning_starts: int = 0, - num_timesteps: int = 0, replay_buffer=None, obs: Optional[np.ndarray] = None, episode_num: int = 0, @@ -537,7 +547,6 @@ class BaseRLModel(ABC): :param deterministic: (bool) :param callback: (BaseCallback) :param learning_starts: (int) - :param num_timesteps: (int) :param replay_buffer: (ReplayBuffer) :param obs: (np.ndarray) :param episode_num: (int) @@ -583,7 +592,7 @@ class BaseRLModel(ABC): # Select action randomly or according to policy # TODO: use action from policy when using SDE during the warmup phase? # if num_timesteps < learning_starts and not self.use_sde: - if num_timesteps < learning_starts: + if self.num_timesteps < learning_starts: # Warmup phase unscaled_action = np.array([self.action_space.sample()]) else: @@ -642,7 +651,7 @@ class BaseRLModel(ABC): if self._vec_normalize_env is not None: obs_ = new_obs_ - num_timesteps += 1 + self.num_timesteps += 1 episode_timesteps += 1 total_steps += 1 if 0 < n_steps <= total_steps: @@ -658,8 +667,8 @@ class BaseRLModel(ABC): # Display training infos if self.verbose >= 1 and log_interval is not None and ( - episode_num + total_episodes) % log_interval == 0: - fps = int(num_timesteps / (time.time() - self.start_time)) + episode_num + total_episodes) % log_interval == 0: + fps = int(self.num_timesteps / (time.time() - self.start_time)) logger.logkv("episodes", episode_num + total_episodes) if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer])) @@ -667,7 +676,7 @@ class BaseRLModel(ABC): # logger.logkv("n_updates", n_updates) logger.logkv("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) - logger.logkv("total timesteps", num_timesteps) + logger.logkv("total timesteps", self.num_timesteps) if self.use_sde: logger.logkv("std", (self.actor.get_std()).mean().item()) logger.dumpkvs() @@ -701,7 +710,7 @@ class BaseRLModel(ABC): @staticmethod def _save_to_file_zip(save_path: str, data: Dict[str, Any] = None, - params: TensorDict = None, opt_params: OptimizerStateDict = None) -> None: + params: TensorDict = None, opt_params: OptimizerStateDict = None) -> None: """ Save model to a zip archive. @@ -776,27 +785,3 @@ class BaseRLModel(ABC): params_to_save = self.get_policy_parameters() opt_params_to_save = self.get_opt_parameters() self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save) - - def _eval_policy(self, eval_freq: int, eval_env: GymEnv, n_eval_episodes: int, - timesteps_since_eval: int, render: bool = False, deterministic: bool = True) -> int: - """ - Evaluate the current policy on a test environment. - - :param eval_freq: Evaluate the agent every `eval_freq` timesteps (this may vary a little) - :param n_eval_episodes: Number of episode to evaluate the agent - :parma timesteps_since_eval: Number of timesteps since last evaluation - :param deterministic: Whether to use deterministic or stochastic actions - :param render: Whether to render the eval env or not - :return: Number of timesteps since last evaluation - """ - if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: - timesteps_since_eval %= eval_freq - # Synchronise the normalization stats if needed - sync_envs_normalization(self.env, eval_env) - mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes, - render=render, deterministic=deterministic) - if self.verbose > 0: - print(f"Eval num_timesteps={self.num_timesteps}, " - f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}") - print(f"FPS: {self.num_timesteps / (time.time() - self.start_time):.2f}") - return timesteps_since_eval diff --git a/torchy_baselines/common/callbacks.py b/torchy_baselines/common/callbacks.py index 061d6e5..89b1e22 100644 --- a/torchy_baselines/common/callbacks.py +++ b/torchy_baselines/common/callbacks.py @@ -1,15 +1,18 @@ import os from abc import ABC, abstractmethod +import typing from typing import Union, List, Dict, Any, Optional import gym import numpy as np -from torchy_baselines.common.base_class import BaseRLModel # pytype: disable=pyi-error from torchy_baselines.common.vec_env import VecEnv, sync_envs_normalization from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.logger import Logger +if typing.TYPE_CHECKING: + from torchy_baselines.common.base_class import BaseRLModel # pytype: disable=pyi-error + class BaseCallback(ABC): """ @@ -31,7 +34,8 @@ class BaseCallback(ABC): # to have access to the parent object self.parent = None # type: Optional[BaseCallback] - def init_callback(self, model: BaseRLModel) -> None: + # Type hint as string to avoid circular import + def init_callback(self, model: 'BaseRLModel') -> None: """ Initialize the callback by saving references to the RL model and the training environment for convenience. @@ -105,18 +109,23 @@ class EventCallback(BaseCallback): if callback is not None: self.callback.parent = self - def init_callback(self, model: BaseRLModel) -> None: + def init_callback(self, model: 'BaseRLModel') -> None: super(EventCallback, self).init_callback(model) - self.callback.init_callback(self.model) + if self.callback is not None: + self.callback.init_callback(self.model) def _on_training_start(self) -> None: - self.callback.on_training_start(self.locals, self.globals) + if self.callback is not None: + self.callback.on_training_start(self.locals, self.globals) def _on_event(self) -> bool: if self.callback is not None: return self.callback() return True + def _on_step(self) -> bool: + return True + class CallbackList(BaseCallback): def __init__(self, callbacks: List[BaseCallback]): @@ -179,7 +188,7 @@ class ConvertCallback(BaseCallback): """ Convert functional callback (old-style) to object. - :param on_step: (callable) + :param callback: (callable) :param verbose: (int) """ def __init__(self, callback, verbose=0): @@ -207,6 +216,7 @@ class EvalCallback(EventCallback): according to performance on the eval env will be saved. :param deterministic: (bool) Whether the evaluation should use a stochastic or deterministic actions. + :param deterministic: (bool) Whether to render or not the environment during evaluation :param verbose: (int) """ def __init__(self, eval_env: Union[gym.Env, VecEnv], @@ -216,12 +226,15 @@ class EvalCallback(EventCallback): log_path: str = None, best_model_save_path: str = None, deterministic: bool = True, + render: bool = False, verbose: int = 1): super(EvalCallback, self).__init__(callback_on_new_best, verbose=verbose) self.n_eval_episodes = n_eval_episodes self.eval_freq = eval_freq self.best_mean_reward = -np.inf self.deterministic = deterministic + self.render = render + if isinstance(eval_env, VecEnv): assert eval_env.num_envs == 1, "You must pass only one environment for evaluation" @@ -230,6 +243,7 @@ class EvalCallback(EventCallback): self.log_path = log_path self.evaluations_results = [] self.evaluations_timesteps = [] + self.evaluations_length = [] def _init_callback(self): # Does not work when eval_env is a gym.Env and training_env is a VecEnv @@ -244,22 +258,30 @@ class EvalCallback(EventCallback): def _on_step(self) -> bool: - if self.n_calls % self.eval_freq == 0: + if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0: # Sync training and eval env if there is VecNormalize sync_envs_normalization(self.training_env, self.eval_env) - episode_rewards, _ = evaluate_policy(self.model, self.eval_env, n_eval_episodes=self.n_eval_episodes, - deterministic=self.deterministic, return_episode_rewards=True) + episode_rewards, episode_lengths = evaluate_policy(self.model, self.eval_env, + n_eval_episodes=self.n_eval_episodes, + render=self.render, + deterministic=self.deterministic, + return_episode_rewards=True) if self.log_path is not None: self.evaluations_timesteps.append(self.num_timesteps) self.evaluations_results.append(episode_rewards) - np.savez(self.log_path, timesteps=self.evaluations_timesteps, results=self.evaluations_results) + self.evaluations_length.append(episode_lengths) + np.savez(self.log_path, timesteps=self.evaluations_timesteps, + results=self.evaluations_results, ep_lengths=self.evaluations_length) mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards) + mean_ep_length, std_ep_length = np.mean(episode_lengths), np.std(episode_lengths) + if self.verbose > 0: print(f"Eval num_timesteps={self.num_timesteps}, " f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}") + print(f"Episode length: {mean_ep_length:.2f} +/- {std_ep_length:.2f}") if mean_reward > self.best_mean_reward: if self.verbose > 0: diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py index 0133621..ff9fa13 100644 --- a/torchy_baselines/common/evaluation.py +++ b/torchy_baselines/common/evaluation.py @@ -24,31 +24,33 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, :param return_episode_rewards: (bool) If True, a list of reward per episode will be returned instead of the mean. :return: (float, float) Mean reward per episode, std of reward per episode - returns ([float], int) when `return_episode_rewards` is True + returns ([float], [int]) when `return_episode_rewards` is True """ if isinstance(env, VecEnv): assert env.num_envs == 1, "You must pass only one environment when using this function" - episode_rewards, n_steps = [], 0 + episode_rewards, episode_lengths = [], [] for _ in range(n_eval_episodes): obs = env.reset() done = False episode_reward = 0.0 + episode_length = 0 while not done: action = model.predict(obs, deterministic=deterministic) obs, reward, done, _info = env.step(action) episode_reward += reward if callback is not None: callback(locals(), globals()) - n_steps += 1 + episode_length += 1 if render: env.render() episode_rewards.append(episode_reward) + episode_lengths.append(episode_length) mean_reward = np.mean(episode_rewards) std_reward = np.std(episode_rewards) if reward_threshold is not None: assert mean_reward > reward_threshold, (f'Mean reward below threshold: ' '{mean_reward:.2f} < {reward_threshold:.2f}') if return_episode_rewards: - return episode_rewards, n_steps + return episode_rewards, episode_lengths return mean_reward, std_reward diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 77c19b4..39d165e 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -193,6 +193,8 @@ class PPO(BaseRLModel): self._update_info_buffer(infos) n_steps += 1 + self.num_timesteps += env.num_envs + if isinstance(self.action_space, gym.spaces.Discrete): # Reshape in case of discrete action actions = actions.reshape(-1, 1) @@ -284,9 +286,11 @@ class PPO(BaseRLModel): logger.logkv("std", th.exp(self.policy.log_std).mean().item()) def learn(self, total_timesteps, callback=None, log_interval=1, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", + eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, iteration, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) + iteration = 0 if self.tensorboard_log is not None and SummaryWriter is not None: self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name)) @@ -295,15 +299,15 @@ class PPO(BaseRLModel): while self.num_timesteps < total_timesteps: - obs, continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps, + obs, continue_training = self.collect_rollouts(self.env, callback, + self.rollout_buffer, + n_rollout_steps=self.n_steps, obs=obs) if continue_training is False: break iteration += 1 - self.num_timesteps += self.n_steps * self.n_envs - timesteps_since_eval += self.n_steps * self.n_envs self._update_current_progress(self.num_timesteps, total_timesteps) # Display training infos @@ -320,9 +324,6 @@ class PPO(BaseRLModel): self.train(self.n_epochs, batch_size=self.batch_size) - # Evaluate the agent - timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes, - timesteps_since_eval, deterministic=True) # For tensorboard integration # if self.tb_writer is not None: # self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 3e2b96f..ec24d86 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -257,9 +257,9 @@ class SAC(BaseRLModel): def learn(self, total_timesteps, callback=None, log_interval=4, eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC", - reset_num_timesteps=True): + eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) callback.on_training_start(locals(), globals()) @@ -268,7 +268,6 @@ class SAC(BaseRLModel): n_steps=self.train_freq, action_noise=self.action_noise, deterministic=False, callback=callback, learning_starts=self.learning_starts, - num_timesteps=self.num_timesteps, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, log_interval=log_interval) @@ -278,9 +277,7 @@ class SAC(BaseRLModel): if continue_training is False: break - self.num_timesteps += episode_timesteps episode_num += n_episodes - timesteps_since_eval += episode_timesteps self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: @@ -288,9 +285,6 @@ class SAC(BaseRLModel): self.train(gradient_steps, batch_size=self.batch_size) - timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes, - timesteps_since_eval, deterministic=True) - callback.on_training_end() return self diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 4c0075d..9465d96 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -249,9 +249,10 @@ class TD3(BaseRLModel): del self.rollout_data def learn(self, total_timesteps, callback=None, log_interval=4, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, + tb_log_name="TD3", eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) callback.on_training_start(locals(), globals()) @@ -261,7 +262,6 @@ class TD3(BaseRLModel): n_steps=self.train_freq, action_noise=self.action_noise, deterministic=False, callback=callback, learning_starts=self.learning_starts, - num_timesteps=self.num_timesteps, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, log_interval=log_interval) @@ -272,8 +272,6 @@ class TD3(BaseRLModel): break episode_num += n_episodes - self.num_timesteps += episode_timesteps - timesteps_since_eval += episode_timesteps self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: @@ -290,10 +288,6 @@ class TD3(BaseRLModel): gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) - # Evaluate the agent - timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes, - timesteps_since_eval, deterministic=True) - callback.on_training_end() return self