diff --git a/README.md b/README.md index 74b55c6..cc9f015 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,9 @@ TODO: - SDE: reduce the number of parameters (only n_features instead of n_features x n_actions) for A2C (done for TD3) - SDE: learn the feature extractor? +- Refactor: buffer with numpy array instead of pytorch +- Refactor: remove duplicated code for evaluation +- plotting? -> zoo Later: - get_parameters / set_parameters diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 36ab9aa..68a1c44 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -25,6 +25,7 @@ class A2C(PPO): (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel) :param gamma: (float) Discount factor :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator + Equivalent to classic advantage when set to 1. :param ent_coef: (float) Entropy coefficient for the loss calculation :param vf_coef: (float) Value function coefficient for the loss calculation :param max_grad_norm: (float) The maximum value for the gradient clipping diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 024ac25..0ada89c 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -186,10 +186,6 @@ class CEMRL(TD3): actor_steps += episode_timesteps self.fitnesses.append(episode_reward) - if self.verbose > 1: - print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( - self.num_timesteps, episode_num, episode_timesteps, episode_reward)) - self._update_current_progress(self.num_timesteps, total_timesteps) self.es.tell(self.es_params, self.fitnesses) timesteps_since_eval += actor_steps diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 36cd0c5..34da520 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -20,13 +20,19 @@ class BaseRLModel(object): :param policy: (BasePolicy) Policy object :param env: (Gym environment) The environment to learn from (if registered in Gym, can be str. Can be None for loading trained models) - :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug :param policy_base: (BasePolicy) the base policy used by this method + :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation + :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug :param device: (str or th.device) Device on which the code should. By default, it will try to use a Cuda compatible device and fallback to cpu if it is not possible. + :param support_multi_env: (bool) Whether the algorithm supports training + with multiple environments (as in A2C) + :param create_eval_env: (bool) Whether to create a second environment that will be + used for evaluating the agent periodically. (Only available when passing string for the environment) :param monitor_wrapper: (bool) When creating an environment, whether to wrap it or not in a Monitor wrapper. + :param seed: (int) Seed for the pseudo random generators """ __metaclass__ = ABCMeta @@ -65,6 +71,7 @@ class BaseRLModel(object): # this is used to update the learning rate self._current_progress = 1 + # Create and wrap the env if needed if env is not None: if isinstance(env, str): if create_eval_env: @@ -94,6 +101,12 @@ class BaseRLModel(object): " environment.") def _get_eval_env(self, eval_env): + """ + Return the environment that will be used for evaluation. + + :param eval_env: (gym.Env or VecEnv) + :return: (VecEnv) + """ if eval_env is None: eval_env = self.eval_env @@ -133,7 +146,7 @@ class BaseRLModel(object): """ Compute current progress (from 1 to 0) - :param num_timesteps: (int) + :param num_timesteps: (int) current number of timesteps :param total_timesteps: (int) """ self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps) @@ -169,7 +182,7 @@ class BaseRLModel(object): """ returns the current environment (can be None if not defined) - :return: (Gym Environment) The current environment + :return: (gym.Env) The current environment """ return self.env @@ -177,7 +190,7 @@ class BaseRLModel(object): """ Checks the validity of the environment, and if it is coherent, set it as the current environment. - :param env: (Gym Environment) The environment for learning a policy + :param env: (gym.Env) The environment for learning a policy """ pass @@ -311,21 +324,33 @@ class BaseRLModel(object): self.eval_env.seed(seed) def _setup_learn(self, eval_env): + """ + Initialize different variables needed for training. + + :param eval_env: (gym.Env or VecEnv) + :return: (int, int, [float], np.ndarray, VecEnv) + """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) + if self.action_noise is not None: self.action_noise.reset() + 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() return timesteps_since_eval, episode_num, evaluations, obs, eval_env def _update_info_buffer(self, infos): """ - Retrieve reward and episode length if using Monitor wrapper. + Retrieve reward and episode length and update the buffer + if using Monitor wrapper. + :param infos: ([dict]) """ for info in infos: @@ -338,7 +363,23 @@ class BaseRLModel(object): learning_starts=0, num_timesteps=0, replay_buffer=None, obs=None, episode_num=0, log_interval=None): + """ + Collect rollout using the current policy (and possibly fill the replay buffer) + TODO: move this method to off-policy base class. + :param env: (VecEnv) + :param n_episodes: (int) + :param n_steps: (int) + :param action_noise: (ActionNoise) + :param deterministic: (bool) + :param callback: (callable) + :param learning_starts: (int) + :param num_timesteps: (int) + :param replay_buffer: (ReplayBuffer) + :param obs: (np.ndarray) + :param episode_num: (int) + :param log_interval: (int) + """ episode_rewards = [] total_timesteps = [] total_steps, total_episodes = 0, 0 @@ -432,12 +473,10 @@ class BaseRLModel(object): episode_num + total_episodes) % log_interval == 0: fps = int(num_timesteps / (time.time() - self.start_time)) logger.logkv("episodes", episode_num + total_episodes) - # logger.logkv("mean 100 episode reward", mean_reward) 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])) logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer])) # logger.logkv("n_updates", n_updates) - # logger.logkv("current_lr", current_lr) logger.logkv("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv("total timesteps", num_timesteps) diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 2c3d1fb..ae31a1f 100644 --- a/torchy_baselines/common/buffers.py +++ b/torchy_baselines/common/buffers.py @@ -5,6 +5,15 @@ from torchy_baselines.common.vec_env import unwrap_vec_normalize class BaseBuffer(object): + """ + Base class that represent a buffer (rollout or replay) + + :param buffer_size: (int) Max number of element in the buffer + :param obs_dim: (int) Dimension of the observation + :param action_dim: (int) Dimension of the action space + :param device: (th.device) + :param n_envs: (int) Number of parallel environments + """ def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): super(BaseBuffer, self).__init__() self.buffer_size = buffer_size @@ -31,27 +40,43 @@ class BaseBuffer(object): return tensor.transpose(0, 1).reshape(shape[0] * shape[1], *shape[2:]) def size(self): + """ + :return: (int) The current size of the buffer + """ if self.full: return self.buffer_size return self.pos - def get_pos(self): - return self.pos - def add(self, *args, **kwargs): + """ + Add elements to the buffer. + """ raise NotImplementedError() def reset(self): + """ + Reset the buffer. + """ self.pos = 0 self.full = False def sample(self, batch_size, env=None): + """ + :param batch_size: (int) Number of element to sample + :param env: (VecNormalize) [Optional] associated gym VecEnv + to normalize the observations/rewards when sampling + """ upper_bound = self.buffer_size if self.full else self.pos batch_inds = th.LongTensor( np.random.randint(0, upper_bound, size=batch_size)) return self._get_samples(batch_inds, env=env) def _get_samples(self, batch_inds, env=None): + """ + :param batch_inds: (th.Tensor) + :param env: (gym.Env) + :return: ([th.Tensor]) + """ raise NotImplementedError() def _normalize_obs(self, obs, env=None): @@ -68,9 +93,15 @@ class BaseBuffer(object): class ReplayBuffer(BaseBuffer): """ - Taken from https://github.com/apourchot/CEM-RL - """ + Replay buffer used in off-policy algorithms like SAC/TD3. + Adapted from from https://github.com/apourchot/CEM-RL + :param buffer_size: (int) Max number of element in the buffer + :param obs_dim: (int) Dimension of the observation + :param action_dim: (int) Dimension of the action space + :param device: (th.device) + :param n_envs: (int) Number of parallel environments + """ def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs) @@ -103,6 +134,18 @@ class ReplayBuffer(BaseBuffer): class RolloutBuffer(BaseBuffer): + """ + Rollout buffer used in on-policy algorithms like A2C/PPO. + + :param buffer_size: (int) Max number of element in the buffer + :param obs_dim: (int) Dimension of the observation + :param action_dim: (int) Dimension of the action space + :param device: (th.device) + :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator + Equivalent to classic advantage when set to 1. + :param gamma: (float) Discount factor + :param n_envs: (int) Number of parallel environments + """ def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', gae_lambda=1, gamma=0.99, n_envs=1): super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs) @@ -128,7 +171,10 @@ class RolloutBuffer(BaseBuffer): def compute_returns_and_advantage(self, last_value, dones=False, use_gae=True): """ - From Stable-Baselines PPO2 + Post-processing step: compute the returns (sum of discounted rewards) + and advantage (A(s) = R - V(S)). + Adapted from Stable-Baselines PPO2. + :param last_value: (th.Tensor) :param dones: ([bool]) :param use_gae: (bool) Whether to use Generalized Advantage Estimation @@ -164,6 +210,16 @@ class RolloutBuffer(BaseBuffer): self.advantages = self.returns - self.values def add(self, obs, action, reward, done, value, log_prob): + """ + :param obs: (np.ndarray) Observation + :param action: (np.ndarray) Action + :param reward: (np.ndarray) + :param done: (np.ndarray) End of episode signal. + :param value: (th.Tensor) estimated value of the current state + following the current policy. + :param log_prob: (th.Tensor) log probability of the action + following the current policy. + """ if len(log_prob.shape) == 0: # Reshape 0-d tensor to avoid error log_prob = log_prob.reshape(-1, 1) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 5def28f..c14623c 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -257,9 +257,6 @@ class SAC(BaseRLModel): self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: - if self.verbose > 1: - print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( - self.num_timesteps, episode_num, episode_timesteps, episode_reward)) gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps self.train(gradient_steps, batch_size=self.batch_size) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index ddd7e9d..2e312ac 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -268,9 +268,6 @@ class TD3(BaseRLModel): self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: - if self.verbose > 1: - print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( - self.num_timesteps, episode_num, episode_timesteps, episode_reward)) if self.use_sde: if self.sde_log_std_scheduler is not None: