Merge pull request #1 from Antonin-Raffin/feat/a2c

Add A2C algorithm

README.md

@@ -8,6 +8,7 @@ PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines
 
 ## Implemented Algorithms
 
+- A2C
 - CEM-RL (with TD3)
 - PPO
 - SAC
@@ -18,11 +19,8 @@ PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines
 
 TODO:
 - save/load
-- predict
-- flexible mlp
-- logger
-- better monitor wrapper?
-- A2C
+- better predict
+- complete logger
 
 Later:
 - get_parameters / set_parameters

tests/test_run.py

@@ -3,7 +3,7 @@ import os
 import pytest
 import numpy as np
 
-from torchy_baselines import TD3, CEMRL, PPO, SAC
+from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
 from torchy_baselines.common.noise import NormalActionNoise
 
 
@@ -28,9 +28,10 @@ def test_cemrl():
     os.remove("test_save.pth")
 
 
+@pytest.mark.parametrize("model_class", [A2C, PPO])
 @pytest.mark.parametrize("env_id", ['CartPole-v1', 'Pendulum-v0'])
-def test_ppo(env_id):
-    model = PPO('MlpPolicy', env_id, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
+def test_onpolicy(model_class, env_id):
+    model = model_class('MlpPolicy', env_id, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
     model.learn(total_timesteps=1000, eval_freq=500)
     # model.save("test_save")
     # model.load("test_save")

torchy_baselines/__init__.py

@@ -1,6 +1,7 @@
+from torchy_baselines.a2c import A2C
 from torchy_baselines.cem_rl import CEMRL
 from torchy_baselines.ppo import PPO
 from torchy_baselines.sac import SAC
 from torchy_baselines.td3 import TD3
 
-__version__ = "0.0.4"
+__version__ = "0.0.5a"

torchy_baselines/a2c/__init__.py

@@ -0,0 +1,2 @@
+from torchy_baselines.a2c.a2c import A2C
+from torchy_baselines.ppo.policies import MlpPolicy

torchy_baselines/a2c/a2c.py

@@ -0,0 +1,126 @@
+from gym import spaces
+import torch as th
+import torch.nn.functional as F
+
+from torchy_baselines.common.utils import explained_variance
+from torchy_baselines.ppo.ppo import PPO
+from torchy_baselines.ppo.policies import PPOPolicy
+
+
+class A2C(PPO):
+    """
+    Advantage Actor Critic (A2C)
+
+    Paper: https://arxiv.org/abs/1602.01783
+    Code: This implementation borrows code from https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail and
+    and Stable Baselines (https://github.com/hill-a/stable-baselines)
+
+    Introduction to A2C: https://hackernoon.com/intuitive-rl-intro-to-advantage-actor-critic-a2c-4ff545978752
+
+    :param policy: (PPOPolicy or str) The policy model to use (MlpPolicy, CnnPolicy, ...)
+    :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str)
+    :param learning_rate: (float or callable) The learning rate, it can be a function
+    :param n_steps: (int) The number of steps to run for each environment per update
+        (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
+    :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
+    :param rms_prop_eps: (float) RMSProp epsilon. It stabilizes square root computation in denominator
+        of RMSProp update
+    :param use_rms_prop: (bool) Whether to use RMSprop (default) or Adam as optimizer
+    :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
+        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
+    :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 tensorflow debug
+    :param seed: (int) Seed for the pseudo random generators
+    :param device: (str or th.device) Device (cpu, cuda, ...) on which the code should be run.
+        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,
+                 normalize_advantage=False, tensorboard_log=None, create_eval_env=False,
+                 policy_kwargs=None, verbose=0, seed=0, device='auto',
+                 _init_setup_model=True):
+
+        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,
+                                  tensorboard_log=tensorboard_log, policy_kwargs=policy_kwargs,
+                                  verbose=verbose, device=device, create_eval_env=create_eval_env,
+                                  seed=seed, _init_setup_model=False)
+
+        self.normalize_advantage = normalize_advantage
+        self.rms_prop_eps = rms_prop_eps
+        self.use_rms_prop = use_rms_prop
+
+        if _init_setup_model:
+            self._setup_model()
+
+    def _setup_model(self):
+        super(A2C, self)._setup_model()
+        if self.use_rms_prop:
+            self.policy.optimizer = th.optim.RMSprop(self.policy.parameters(),
+                                                     lr=self.learning_rate(1), alpha=0.99,
+                                                     eps=self.rms_prop_eps, weight_decay=0)
+
+    def train(self, gradient_steps, batch_size=None):
+
+        # Update optimizer learning rate
+        self._update_learning_rate(self.policy.optimizer)
+        # A2C with gradient_steps > 1 does not make sense
+        assert gradient_steps == 1
+        # We do not use minibatches for A2C
+        assert batch_size is None
+
+        for rollout_data in self.rollout_buffer.get(batch_size=None):
+            # Unpack
+            obs, action, _, _, advantage, return_batch = rollout_data
+
+            if isinstance(self.action_space, spaces.Discrete):
+                # Convert discrete action for float to long
+                action = action.long().flatten()
+
+            # TODO: avoid second computation of everything because of the gradient
+            values, log_prob, entropy = self.policy.get_policy_stats(obs, action)
+            values = values.flatten()
+
+            # Normalize advantage (not present in the original implementation)
+            if self.normalize_advantage:
+                advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
+
+            policy_loss = -(advantage * log_prob).mean()
+
+            # Value loss using the TD(gae_lambda) target
+            value_loss = F.mse_loss(return_batch, values)
+
+            # Entropy loss favor exploration
+            entropy_loss = -th.mean(entropy)
+
+            loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss
+
+            # Optimization step
+            self.policy.optimizer.zero_grad()
+            loss.backward()
+            # Clip grad norm
+            th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
+            self.policy.optimizer.step()
+            # approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy())
+
+        # print(explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
+        #                          self.rollout_buffer.values.flatten().cpu().numpy()))
+
+    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):
+
+        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)

torchy_baselines/cem_rl/cem_rl.py

@@ -78,7 +78,7 @@ class CEMRL(TD3):
                     # set params
                     self.actor.load_from_vector(self.es_params[i])
                     self.actor_target.load_from_vector(self.es_params[i])
-                    self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=self.learning_rate)
+                    self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=self.learning_rate(self._current_progress))
 
                     # In the paper: 2 * actor_steps // self.n_grad
                     # In the original implementation: actor_steps // self.n_grad
@@ -153,6 +153,7 @@ class CEMRL(TD3):
                     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
         return self

torchy_baselines/common/base_class.py

@@ -7,7 +7,7 @@ import torch as th
 import numpy as np
 
 from torchy_baselines.common.policies import get_policy_from_name
-from torchy_baselines.common.utils import set_random_seed
+from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
 from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv
 from torchy_baselines.common.monitor import Monitor
 from torchy_baselines.common import logger
@@ -57,6 +57,9 @@ class BaseRLModel(object):
         self.replay_buffer = None
         self.seed = seed
         self.action_noise = None
+        # Track the training progress (from 1 to 0)
+        # this is used to update the learning rate
+        self._current_progress = 1
 
         if env is not None:
             if isinstance(env, str):
@@ -112,6 +115,35 @@ class BaseRLModel(object):
         low, high = self.action_space.low, self.action_space.high
         return low + (0.5 * (scaled_action + 1.0) * (high -  low))
 
+    def _setup_learning_rate(self):
+        """Transform to callable if needed."""
+        self.learning_rate = get_schedule_fn(self.learning_rate)
+
+    def _update_current_progress(self, num_timesteps, total_timesteps):
+        """
+        Compute current progress (from 1 to 0)
+
+        :param num_timesteps: (int)
+        :param total_timesteps: (int)
+        """
+        self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps)
+
+    def _update_learning_rate(self, optimizers):
+        """
+        Update the optimizers learning rate using the current learning rate schedule
+        and the current progress (from 1 to 0).
+
+        :param optimizers: ([th.optim.Optimizer] or Optimizer) An optimizer
+            or a list of optimizer.
+        """
+        # Log the current learning rate
+        logger.logkv("learning_rate", self.learning_rate(self._current_progress))
+
+        if not isinstance(optimizers, list):
+            optimizers = [optimizers]
+        for optimizer in optimizers:
+            update_learning_rate(optimizer, self.learning_rate(self._current_progress))
+
     @staticmethod
     def safe_mean(arr):
         """

torchy_baselines/common/buffers.py

@@ -134,6 +134,7 @@ class RolloutBuffer(BaseBuffer):
         if len(log_prob.shape) == 0:
             # Reshape 0-d tensor to avoid error
             log_prob = log_prob.reshape(-1, 1)
+
         self.observations[self.pos] = th.FloatTensor(np.array(obs).copy())
         self.actions[self.pos] = th.FloatTensor(np.array(action).copy())
         self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy())
@@ -144,7 +145,7 @@ class RolloutBuffer(BaseBuffer):
         if self.pos == self.buffer_size:
             self.full = True
 
-    def get(self, batch_size):
+    def get(self, batch_size=None):
         assert self.full
         indices = th.randperm(self.buffer_size * self.n_envs)
         # Prepare the data
@@ -154,6 +155,10 @@ class RolloutBuffer(BaseBuffer):
                 self.__dict__[tensor] = self.swap_and_flatten(self.__dict__[tensor])
             self.generator_ready = True
 
+        # Return everything, don't create minibatches
+        if batch_size is None:
+            batch_size = self.buffer_size * self.n_envs
+
         start_idx = 0
         while start_idx < self.buffer_size * self.n_envs:
             yield self._get_samples(indices[start_idx:start_idx + batch_size])

torchy_baselines/common/utils.py

@@ -38,3 +38,48 @@ def explained_variance(y_pred, y_true):
     assert y_true.ndim == 1 and y_pred.ndim == 1
     var_y = np.var(y_true)
     return np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
+
+
+def update_learning_rate(optimizer, learning_rate):
+    """
+    Update the learning rate for a given optimizer.
+    Useful when doing linear schedule.
+
+    :param optimizer: (th.optim.Optimizer)
+    :param learning_rate: (float)
+    """
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = learning_rate
+
+
+def get_schedule_fn(value_schedule):
+    """
+    Transform (if needed) learning rate and clip range (for PPO)
+    to callable.
+
+    :param value_schedule: (callable or float)
+    :return: (function)
+    """
+    # If the passed schedule is a float
+    # create a constant function
+    if isinstance(value_schedule, (float, int)):
+        # Cast to float to avoid errors
+        value_schedule = constant_fn(float(value_schedule))
+    else:
+        assert callable(value_schedule)
+    return value_schedule
+
+
+def constant_fn(val):
+    """
+    Create a function that returns a constant
+    It is useful for learning rate schedule (to avoid code duplication)
+
+    :param val: (float)
+    :return: (function)
+    """
+
+    def func(_):
+        return val
+
+    return func

torchy_baselines/ppo/policies.py

@@ -100,7 +100,7 @@ class MlpExtractor(nn.Module):
 
 class PPOPolicy(BasePolicy):
     def __init__(self, observation_space, action_space,
-                 learning_rate=1e-3, net_arch=None, device='cpu',
+                 learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.Tanh, adam_epsilon=1e-5, ortho_init=True):
         super(PPOPolicy, self).__init__(observation_space, action_space, device)
         self.obs_dim = self.observation_space.shape[0]
@@ -148,8 +148,7 @@ class PPOPolicy(BasePolicy):
                     self.value_net: 1
                 }[module]
                 module.apply(partial(self.init_weights, gain=gain))
-        # TODO: support linear decay of the learning rate
-        self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate, eps=self.adam_epsilon)
+        self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate(1), eps=self.adam_epsilon)
 
     def forward(self, obs, deterministic=False):
         if not isinstance(obs, th.Tensor):

torchy_baselines/ppo/ppo.py

@@ -16,7 +16,7 @@ import numpy as np
 from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.buffers import RolloutBuffer
-from torchy_baselines.common.utils import explained_variance
+from torchy_baselines.common.utils import explained_variance, get_schedule_fn
 from torchy_baselines.common.vec_env import VecNormalize
 from torchy_baselines.common import logger
 from torchy_baselines.ppo.policies import PPOPolicy
@@ -36,14 +36,16 @@ class PPO(BaseRLModel):
     :param policy: (PPOPolicy or str) The policy model to use (MlpPolicy, CnnPolicy, ...)
     :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str)
     :param learning_rate: (float or callable) The learning rate, it can be a function
+        of the current progress (from 1 to 0)
     :param n_steps: (int) The number of steps to run for each environment per update
         (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
     :param batch_size: (int) Minibatch size
     :param n_epochs: (int) Number of epoch when optimizing the surrogate loss
     :param gamma: (float) Discount factor
     :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
-    :param clip_range: (float or callable) Clipping parameter, it can be a function
-    :param clip_range_vf: (float or callable) Clipping parameter for the value function, it can be a function.
+    :param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress (from 1 to 0).
+    :param clip_range_vf: (float or callable) Clipping parameter for the value function,
+        it can be a function of the current progress (from 1 to 0).
         This is a parameter specific to the OpenAI implementation. If None is passed (default),
         no clipping will be done on the value function.
         IMPORTANT: this clipping depends on the reward scaling.
@@ -84,12 +86,12 @@ class PPO(BaseRLModel):
         self.gamma = gamma
         self.gae_lambda = gae_lambda
         self.clip_range = clip_range
+        self.clip_range_vf = clip_range_vf
         self.ent_coef = ent_coef
         self.vf_coef = vf_coef
         self.max_grad_norm = max_grad_norm
         self.rollout_buffer = None
         self.target_kl = target_kl
-        self.clip_range_vf = clip_range_vf
         self.tensorboard_log = tensorboard_log
         self.tb_writer = None
 
@@ -97,6 +99,7 @@ class PPO(BaseRLModel):
             self._setup_model()
 
     def _setup_model(self):
+        self._setup_learning_rate()
         # TODO: preprocessing: one hot vector for obs discrete
         state_dim = self.observation_space.shape[0]
         if isinstance(self.action_space, spaces.Box):
@@ -116,6 +119,10 @@ class PPO(BaseRLModel):
                                   self.learning_rate, device=self.device, **self.policy_kwargs)
         self.policy = self.policy.to(self.device)
 
+        self.clip_range = get_schedule_fn(self.clip_range)
+        if self.clip_range_vf is not None:
+            self.clip_range_vf = get_schedule_fn(self.clip_range_vf)
+
     def select_action(self, observation):
         # Normally not needed
         observation = np.array(observation)
@@ -158,6 +165,9 @@ class PPO(BaseRLModel):
 
             self._update_info_buffer(infos)
             n_steps += 1
+            if isinstance(self.action_space, gym.spaces.Discrete):
+                # Reshape in case of discrete action
+                actions = actions.reshape(-1, 1)
             rollout_buffer.add(obs, actions, rewards, dones, values, log_probs)
             obs = new_obs
 
@@ -166,6 +176,15 @@ class PPO(BaseRLModel):
         return obs
 
     def train(self, gradient_steps, batch_size=64):
+        # Update optimizer learning rate
+        self._update_learning_rate(self.policy.optimizer)
+        # Compute current clip range
+        clip_range = self.clip_range(self._current_progress)
+        logger.logkv("clip_range", clip_range)
+        if self.clip_range_vf is not None:
+            clip_range_vf = self.clip_range_vf(self._current_progress)
+            logger.logkv("clip_range_vf", clip_range_vf)
+
 
         for gradient_step in range(gradient_steps):
             approx_kl_divs = []
@@ -187,7 +206,7 @@ class PPO(BaseRLModel):
                 ratio = th.exp(log_prob - old_log_prob)
                 # clipped surrogate loss
                 policy_loss_1 = advantage * ratio
-                policy_loss_2 = advantage * th.clamp(ratio, 1 - self.clip_range, 1 + self.clip_range)
+                policy_loss_2 = advantage * th.clamp(ratio, 1 - clip_range, 1 + clip_range)
                 policy_loss = -th.min(policy_loss_1, policy_loss_2).mean()
 
                 if self.clip_range_vf is None:
@@ -196,13 +215,13 @@ class PPO(BaseRLModel):
                 else:
                     # Clip the different between old and new value
                     # NOTE: this depends on the reward scaling
-                    values_pred = old_values + th.clamp(values - old_values, -self.clip_range_vf, self.clip_range_vf)
+                    values_pred = old_values + th.clamp(values - old_values, -clip_range_vf, clip_range_vf)
                 # Value loss using the TD(gae_lambda) target
                 value_loss = F.mse_loss(return_batch, values_pred)
 
 
                 # Entropy loss favor exploration
-                entropy_loss = th.mean(entropy)
+                entropy_loss = -th.mean(entropy)
 
                 loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss
 
@@ -241,6 +260,7 @@ class PPO(BaseRLModel):
             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
             if self.verbose >= 1 and log_interval is not None and iteration % log_interval == 0:

torchy_baselines/sac/policies.py

@@ -63,7 +63,7 @@ class Critic(BaseNetwork):
 
 class SACPolicy(BasePolicy):
     def __init__(self, observation_space, action_space,
-                 learning_rate=3e-4, net_arch=None, device='cpu',
+                 learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.ReLU):
         super(SACPolicy, self).__init__(observation_space, action_space, device)
 
@@ -87,12 +87,12 @@ class SACPolicy(BasePolicy):
 
     def _build(self, learning_rate):
         self.actor = self.make_actor()
-        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate)
+        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate(1))
 
         self.critic = self.make_critic()
         self.critic_target = self.make_critic()
         self.critic_target.load_state_dict(self.critic.state_dict())
-        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate)
+        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
 
     def make_actor(self):
         return Actor(**self.net_args).to(self.device)

torchy_baselines/sac/sac.py

@@ -88,6 +88,7 @@ class SAC(BaseRLModel):
             self._setup_model()
 
     def _setup_model(self):
+        self._setup_learning_rate()
         obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
         if self.seed is not None:
             self.set_random_seed(self.seed)
@@ -114,7 +115,7 @@ class SAC(BaseRLModel):
             # Note: we optimize the log of the entropy coeff which is slightly different from the paper
             # as discussed in https://github.com/rail-berkeley/softlearning/issues/37
             self.log_ent_coef = th.log(th.ones(1, device=self.device) * init_value).requires_grad_(True)
-            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.learning_rate)
+            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.learning_rate(1))
         else:
             # Force conversion to float
             # this will throw an error if a malformed string (different from 'auto')
@@ -152,6 +153,13 @@ class SAC(BaseRLModel):
         return self.unscale_action(self.select_action(observation))
 
     def train(self, gradient_steps, batch_size=64):
+        # Update optimizers learning rate
+        optimizers = [self.actor.optimizer, self.critic.optimizer]
+        if self.ent_coef_optimizer is not None:
+            optimizers += [self.ent_coef_optimizer]
+
+        self._update_learning_rate(optimizers)
+
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
             replay_data = self.replay_buffer.sample(batch_size)
@@ -245,6 +253,7 @@ class SAC(BaseRLModel):
             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:
                 if self.verbose > 1:

torchy_baselines/td3/policies.py

@@ -39,7 +39,7 @@ class Critic(BaseNetwork):
 
 class TD3Policy(BasePolicy):
     def __init__(self, observation_space, action_space,
-                 learning_rate=1e-3, net_arch=None, device='cpu',
+                 learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.ReLU):
         super(TD3Policy, self).__init__(observation_space, action_space, device)
 
@@ -64,12 +64,12 @@ class TD3Policy(BasePolicy):
         self.actor = self.make_actor()
         self.actor_target = self.make_actor()
         self.actor_target.load_state_dict(self.actor.state_dict())
-        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate)
+        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate(1))
 
         self.critic = self.make_critic()
         self.critic_target = self.make_critic()
         self.critic_target.load_state_dict(self.critic.state_dict())
-        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate)
+        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
 
     def make_actor(self):
         return Actor(**self.net_args).to(self.device)

torchy_baselines/td3/td3.py

@@ -75,6 +75,7 @@ class TD3(BaseRLModel):
             self._setup_model()
 
     def _setup_model(self):
+        self._setup_learning_rate()
         obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
         self.set_random_seed(self.seed)
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
@@ -109,6 +110,8 @@ class TD3(BaseRLModel):
         return self.unscale_action(self.select_action(observation))
 
     def train_critic(self, gradient_steps=1, batch_size=100, replay_data=None, tau=0.0):
+        # Update optimizer learning rate
+        self._update_learning_rate(self.critic.optimizer)
 
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
@@ -146,6 +149,8 @@ class TD3(BaseRLModel):
                     target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
 
     def train_actor(self, gradient_steps=1, batch_size=100, tau_actor=0.005, tau_critic=0.005, replay_data=None):
+        # Update optimizer learning rate
+        self._update_learning_rate(self.actor.optimizer)
 
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
@@ -208,6 +213,7 @@ class TD3(BaseRLModel):
             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:
                 if self.verbose > 1: