From 1d6f9bf1006e6e122637b0ae11ad91bc8dfe0025 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 17 Dec 2019 11:47:21 +0100 Subject: [PATCH] Add sample freq for SDE --- torchy_baselines/a2c/a2c.py | 8 +++++--- torchy_baselines/common/base_class.py | 10 +++++++++- torchy_baselines/ppo/policies.py | 2 +- torchy_baselines/ppo/ppo.py | 15 ++++++++++----- torchy_baselines/sac/sac.py | 7 +++++-- torchy_baselines/td3/td3.py | 6 ++++-- 6 files changed, 34 insertions(+), 14 deletions(-) diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 33eae78..6fa667c 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -33,6 +33,8 @@ class A2C(PPO): :param use_rms_prop: (bool) Whether to use RMSprop (default) or Adam as optimizer :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) instead of action noise exploration (default: False) + :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE + Default: -1 (only sample at the beginning of the rollout) :param normalize_advantage: (bool) Whether to normalize or not the advantage :param tensorboard_log: (str) the log location for tensorboard (if None, no logging) :param create_eval_env: (bool) Whether to create a second environment that will be @@ -44,11 +46,10 @@ class A2C(PPO): Setting it to auto, the code will be run on the GPU if possible. :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ - def __init__(self, policy, env, learning_rate=7e-4, n_steps=5, gamma=0.99, gae_lambda=1.0, ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5, - rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False, + rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False, sde_sample_freq=-1, normalize_advantage=False, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): @@ -56,7 +57,8 @@ class A2C(PPO): super(A2C, self).__init__(policy, env, learning_rate=learning_rate, n_steps=n_steps, batch_size=None, n_epochs=1, gamma=gamma, gae_lambda=gae_lambda, ent_coef=ent_coef, - vf_coef=vf_coef, max_grad_norm=max_grad_norm, use_sde=use_sde, + vf_coef=vf_coef, max_grad_norm=max_grad_norm, + use_sde=use_sde, sde_sample_freq=sde_sample_freq, tensorboard_log=tensorboard_log, policy_kwargs=policy_kwargs, verbose=verbose, device=device, create_eval_env=create_eval_env, seed=seed, _init_setup_model=False) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 576bea2..9073296 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -39,12 +39,15 @@ class BaseRLModel(object): :param seed: (int) Seed for the pseudo random generators :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) instead of action noise exploration (default: False) + :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE + Default: -1 (only sample at the beginning of the rollout) """ __metaclass__ = ABCMeta def __init__(self, policy, env, policy_base, policy_kwargs=None, verbose=0, device='auto', support_multi_env=False, - create_eval_env=False, monitor_wrapper=True, seed=None, use_sde=False): + create_eval_env=False, monitor_wrapper=True, seed=None, + use_sde=False, sde_sample_freq=-1): if isinstance(policy, str) and policy_base is not None: self.policy_class = get_policy_from_name(policy_base, policy) else: @@ -74,6 +77,7 @@ class BaseRLModel(object): self.rollout_data = None self.on_policy_exploration = False self.use_sde = use_sde + self.sde_sample_freq = sde_sample_freq # Track the training progress (from 1 to 0) # this is used to update the learning rate self._current_progress = 1 @@ -516,6 +520,10 @@ class BaseRLModel(object): episode_reward, episode_timesteps = 0.0, 0 while not done: + if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0: + # Sample a new noise matrix + self.actor.reset_noise() + # Select action randomly or according to policy if num_timesteps < learning_starts: # Warmup phase diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 19e20a5..8c4df3f 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -76,7 +76,7 @@ class PPOPolicy(BasePolicy): self._build(learning_rate) - def reset_noise_net(self, n_envs=1): + def reset_noise(self, n_envs=1): """ Sample new weights for the exploration matrix. diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 113e050..bc2c7d8 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -55,6 +55,8 @@ class PPO(BaseRLModel): :param max_grad_norm: (float) The maximum value for the gradient clipping :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) instead of action noise exploration (default: False) + :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE + Default: -1 (only sample at the beginning of the rollout) :param target_kl: (float) Limit the KL divergence between updates, because the clipping is not enough to prevent large update see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213) @@ -69,17 +71,17 @@ class PPO(BaseRLModel): Setting it to auto, the code will be run on the GPU if possible. :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ - def __init__(self, policy, env, learning_rate=3e-4, n_steps=2048, batch_size=64, n_epochs=10, gamma=0.99, gae_lambda=0.95, clip_range=0.2, clip_range_vf=None, - ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5, use_sde=False, + ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5, + use_sde=False, sde_sample_freq=-1, target_kl=None, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs=policy_kwargs, - verbose=verbose, device=device, + verbose=verbose, device=device, use_sde=use_sde, sde_sample_freq=sde_sample_freq, create_eval_env=create_eval_env, support_multi_env=True, seed=seed) self.learning_rate = learning_rate @@ -97,7 +99,6 @@ class PPO(BaseRLModel): self.target_kl = target_kl self.tensorboard_log = tensorboard_log self.tb_writer = None - self.use_sde = use_sde if _init_setup_model: self._setup_model() @@ -158,9 +159,13 @@ class PPO(BaseRLModel): # Sample new weights for the state dependent exploration # TODO: ensure episodic setting? if self.use_sde: - self.policy.reset_noise_net(env.num_envs) + self.policy.reset_noise(env.num_envs) while n_steps < n_rollout_steps: + if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0: + # Sample a new noise matrix + self.policy.reset_noise(env.num_envs) + with th.no_grad(): actions, values, log_probs = self.policy.forward(obs) actions = actions.cpu().numpy() diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index ee5ee4c..bc24aa3 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -48,6 +48,8 @@ class SAC(BaseRLModel): :param gamma: (float) the discount factor :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) instead of action noise exploration (default: False) + :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE + Default: -1 (only sample at the beginning of the rollout) :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 policy_kwargs: (dict) additional arguments to be passed to the policy on creation @@ -63,12 +65,13 @@ class SAC(BaseRLModel): tau=0.005, ent_coef='auto', target_update_interval=1, train_freq=1, gradient_steps=1, n_episodes_rollout=-1, target_entropy='auto', action_noise=None, max_grad_norm=None, - gamma=0.99, use_sde=False, tensorboard_log=None, create_eval_env=False, + gamma=0.99, use_sde=False, sde_sample_freq=-1, + tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): super(SAC, self).__init__(policy, env, SACPolicy, policy_kwargs, verbose, device, - create_eval_env=create_eval_env, seed=seed, use_sde=use_sde) + create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq) self.learning_rate = learning_rate self.target_entropy = target_entropy diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 06eaff5..3e1dfb7 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -40,6 +40,8 @@ class TD3(BaseRLModel): :param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise. :param use_sde: (bool) Whether to use State Dependent Exploration (SDE) instead of action noise exploration (default: False) + :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE + Default: -1 (only sample at the beginning of the rollout) :param sde_max_grad_norm: (float) :param sde_ent_coef: (float) :param sde_log_std_scheduler: (callable) @@ -57,12 +59,12 @@ class TD3(BaseRLModel): policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5, - use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None, + use_sde=False, sde_sample_freq=-1, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, seed=0, device='auto', _init_setup_model=True): super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device, - create_eval_env=create_eval_env, seed=seed, use_sde=use_sde) + create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq) self.buffer_size = buffer_size self.learning_rate = learning_rate