diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e614819..4bc5907 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -27,6 +27,7 @@ Deprecations: Others: ^^^^^^^ - Add more issue templates +- Add signatures to callable type annotations (@erniejunior) - Improve error message in ``NatureCNN`` Documentation: diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index ee25ee2..42400c5 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/stable_baselines3/a2c/a2c.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, Optional, Type, Union +from typing import Any, Dict, Optional, Type, Union import torch as th from gym import spaces @@ -7,7 +7,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm from stable_baselines3.common.policies import ActorCriticPolicy -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import explained_variance @@ -55,7 +55,7 @@ class A2C(OnPolicyAlgorithm): self, policy: Union[str, Type[ActorCriticPolicy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable] = 7e-4, + learning_rate: Union[float, LearningRateSchedule] = 7e-4, n_steps: int = 5, gamma: float = 0.99, gae_lambda: float = 1.0, diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index f33b36f..58bdf88 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -5,7 +5,7 @@ import pathlib import time from abc import ABC, abstractmethod from collections import deque -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union +from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union import gym import numpy as np @@ -18,7 +18,7 @@ from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.policies import BasePolicy, get_policy_from_name from stable_baselines3.common.preprocessing import is_image_space, is_image_space_channels_first from stable_baselines3.common.save_util import load_from_zip_file, recursive_getattr, recursive_setattr, save_to_zip_file -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import ( check_for_correct_spaces, get_device, @@ -89,7 +89,7 @@ class BaseAlgorithm(ABC): policy: Type[BasePolicy], env: Union[GymEnv, str, None], policy_base: Type[BasePolicy], - learning_rate: Union[float, Callable], + learning_rate: Union[float, LearningRateSchedule], policy_kwargs: Dict[str, Any] = None, tensorboard_log: Optional[str] = None, verbose: int = 0, @@ -129,7 +129,7 @@ class BaseAlgorithm(ABC): self.policy = None self.learning_rate = learning_rate self.tensorboard_log = tensorboard_log - self.lr_schedule = None # type: Optional[Callable] + self.lr_schedule = None # type: Optional[LearningRateSchedule] self._last_obs = None # type: Optional[np.ndarray] self._last_dones = None # type: Optional[np.ndarray] # When using VecNormalize: diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 872aed1..27d7748 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -248,7 +248,7 @@ class ConvertCallback(BaseCallback): :param verbose: """ - def __init__(self, callback: Callable, verbose: int = 0): + def __init__(self, callback: Callable[[Dict[str, Any], Dict[str, Any]], bool], verbose: int = 0): super(ConvertCallback, self).__init__(verbose) self.callback = callback diff --git a/stable_baselines3/common/env_util.py b/stable_baselines3/common/env_util.py index 8a9985e..0458b30 100644 --- a/stable_baselines3/common/env_util.py +++ b/stable_baselines3/common/env_util.py @@ -14,7 +14,7 @@ def make_vec_env( seed: Optional[int] = None, start_index: int = 0, monitor_dir: Optional[str] = None, - wrapper_class: Optional[Callable] = None, + wrapper_class: Optional[Callable[[gym.Env], gym.Env]] = None, env_kwargs: Optional[Dict[str, Any]] = None, vec_env_cls: Optional[Type[Union[DummyVecEnv, SubprocVecEnv]]] = None, vec_env_kwargs: Optional[Dict[str, Any]] = None, diff --git a/stable_baselines3/common/evaluation.py b/stable_baselines3/common/evaluation.py index 8b89b00..0ed40e4 100644 --- a/stable_baselines3/common/evaluation.py +++ b/stable_baselines3/common/evaluation.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import gym import numpy as np @@ -13,7 +13,7 @@ def evaluate_policy( n_eval_episodes: int = 10, deterministic: bool = True, render: bool = False, - callback: Optional[Callable] = None, + callback: Optional[Callable[[Dict[str, Any], Dict[str, Any]], None]] = None, reward_threshold: Optional[float] = None, return_episode_rewards: bool = False, ) -> Union[Tuple[float, float], Tuple[List[float], List[int]]]: @@ -28,7 +28,7 @@ def evaluate_policy( :param deterministic: Whether to use deterministic or stochastic actions :param render: Whether to render the environment or not :param callback: callback function to do additional checks, - called after each step. + called after each step. Gets locals() and globals() passed as parameters. :param reward_threshold: Minimum expected reward per episode, this will raise an error if the performance is not met :param return_episode_rewards: If True, a list of reward per episode diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index a536417..37662a6 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -2,7 +2,7 @@ import io import pathlib import time import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Dict, Optional, Tuple, Type, Union import gym import numpy as np @@ -15,7 +15,7 @@ from stable_baselines3.common.callbacks import BaseCallback from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.policies import BasePolicy from stable_baselines3.common.save_util import load_from_pkl, save_to_pkl -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, RolloutReturn +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback, RolloutReturn from stable_baselines3.common.utils import safe_mean from stable_baselines3.common.vec_env import VecEnv @@ -76,7 +76,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): policy: Type[BasePolicy], env: Union[GymEnv, str], policy_base: Type[BasePolicy], - learning_rate: Union[float, Callable], + learning_rate: Union[float, LearningRateSchedule], buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 256, @@ -192,7 +192,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): self, total_timesteps: int, eval_env: Optional[GymEnv], - callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, + callback: MaybeCallback = None, eval_freq: int = 10000, n_eval_episodes: int = 5, log_path: Optional[str] = None, diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 6cd5dbf..527cb9d 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -1,5 +1,5 @@ import time -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Dict, List, Optional, Tuple, Type, Union import gym import numpy as np @@ -10,7 +10,7 @@ from stable_baselines3.common.base_class import BaseAlgorithm from stable_baselines3.common.buffers import RolloutBuffer from stable_baselines3.common.callbacks import BaseCallback from stable_baselines3.common.policies import ActorCriticPolicy -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import safe_mean from stable_baselines3.common.vec_env import VecEnv @@ -52,7 +52,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): self, policy: Union[str, Type[ActorCriticPolicy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable], + learning_rate: Union[float, LearningRateSchedule], n_steps: int, gamma: float, gae_lambda: float, diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index d07f115..64eefd2 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -21,6 +21,7 @@ from stable_baselines3.common.distributions import ( ) from stable_baselines3.common.preprocessing import get_action_dim, is_image_space, preprocess_obs from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, MlpExtractor, NatureCNN, create_mlp +from stable_baselines3.common.type_aliases import LearningRateSchedule from stable_baselines3.common.utils import get_device, is_vectorized_observation from stable_baselines3.common.vec_env import VecTransposeImage from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper @@ -364,7 +365,7 @@ class ActorCriticPolicy(BasePolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable[[float], float], + lr_schedule: LearningRateSchedule, net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, @@ -479,7 +480,7 @@ class ActorCriticPolicy(BasePolicy): self.features_dim, net_arch=self.net_arch, activation_fn=self.activation_fn, device=self.device ) - def _build(self, lr_schedule: Callable[[float], float]) -> None: + def _build(self, lr_schedule: LearningRateSchedule) -> None: """ Create the networks and the optimizer. @@ -662,7 +663,7 @@ class ActorCriticCnnPolicy(ActorCriticPolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, diff --git a/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py b/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py index 758079f..ba70a5f 100644 --- a/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py +++ b/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py @@ -74,7 +74,7 @@ class RMSpropTFLike(Optimizer): group.setdefault("centered", False) @torch.no_grad() - def step(self, closure: Optional[Callable] = None) -> Optional[torch.Tensor]: + def step(self, closure: Optional[Callable[[], None]] = None) -> Optional[torch.Tensor]: """Performs a single optimization step. :param closure: A closure that reevaluates the model diff --git a/stable_baselines3/common/type_aliases.py b/stable_baselines3/common/type_aliases.py index 0d95d00..1577423 100644 --- a/stable_baselines3/common/type_aliases.py +++ b/stable_baselines3/common/type_aliases.py @@ -15,6 +15,7 @@ GymStepReturn = Tuple[GymObs, float, bool, Dict] TensorDict = Dict[str, th.Tensor] OptimizerStateDict = Dict[str, Any] MaybeCallback = Union[None, Callable, List[callbacks.BaseCallback], callbacks.BaseCallback] +LearningRateSchedule = Callable[[float], float] class RolloutBufferSamples(NamedTuple): diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 2f02ed7..5b1043e 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -16,7 +16,7 @@ except ImportError: SummaryWriter = None from stable_baselines3.common import logger -from stable_baselines3.common.type_aliases import GymEnv +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule def set_random_seed(seed: int, using_cuda: bool = False) -> None: @@ -70,7 +70,7 @@ def update_learning_rate(optimizer: th.optim.Optimizer, learning_rate: float) -> param_group["lr"] = learning_rate -def get_schedule_fn(value_schedule: Union[Callable, float]) -> Callable: +def get_schedule_fn(value_schedule: Union[LearningRateSchedule, float, int]) -> LearningRateSchedule: """ Transform (if needed) learning rate and clip range (for PPO) to callable. @@ -88,7 +88,7 @@ def get_schedule_fn(value_schedule: Union[Callable, float]) -> Callable: return value_schedule -def get_linear_fn(start: float, end: float, end_fraction: float) -> Callable: +def get_linear_fn(start: float, end: float, end_fraction: float) -> LearningRateSchedule: """ Create a function that interpolates linearly between start and end between ``progress_remaining`` = 1 and ``progress_remaining`` = ``end_fraction``. @@ -112,7 +112,7 @@ def get_linear_fn(start: float, end: float, end_fraction: float) -> Callable: return func -def constant_fn(val: float) -> Callable: +def constant_fn(val: float) -> LearningRateSchedule: """ Create a function that returns a constant It is useful for learning rate schedule (to avoid code duplication) diff --git a/stable_baselines3/ddpg/ddpg.py b/stable_baselines3/ddpg/ddpg.py index 8acd0f8..d4af6cc 100644 --- a/stable_baselines3/ddpg/ddpg.py +++ b/stable_baselines3/ddpg/ddpg.py @@ -1,10 +1,10 @@ -from typing import Any, Callable, Dict, Optional, Type, Union +from typing import Any, Dict, Optional, Type, Union import torch as th from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.td3.policies import TD3Policy from stable_baselines3.td3.td3 import TD3 @@ -55,7 +55,7 @@ class DDPG(TD3): self, policy: Union[str, Type[TD3Policy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable] = 1e-3, + learning_rate: Union[float, LearningRateSchedule] = 1e-3, buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 100, diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 7a35566..648475e 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Dict, List, Optional, Tuple, Type, Union import numpy as np import torch as th @@ -6,7 +6,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import get_linear_fn, polyak_update from stable_baselines3.dqn.policies import DQNPolicy @@ -59,7 +59,7 @@ class DQN(OffPolicyAlgorithm): self, policy: Union[str, Type[DQNPolicy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable] = 1e-4, + learning_rate: Union[float, LearningRateSchedule] = 1e-4, buffer_size: int = 1000000, learning_starts: int = 50000, batch_size: Optional[int] = 32, diff --git a/stable_baselines3/dqn/policies.py b/stable_baselines3/dqn/policies.py index 7416f02..35e6581 100644 --- a/stable_baselines3/dqn/policies.py +++ b/stable_baselines3/dqn/policies.py @@ -6,6 +6,7 @@ from torch import nn from stable_baselines3.common.policies import BasePolicy, register_policy from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp +from stable_baselines3.common.type_aliases import LearningRateSchedule class QNetwork(BasePolicy): @@ -103,7 +104,7 @@ class DQNPolicy(BasePolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[List[int]] = None, activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, @@ -142,7 +143,7 @@ class DQNPolicy(BasePolicy): self.q_net, self.q_net_target = None, None self._build(lr_schedule) - def _build(self, lr_schedule: Callable) -> None: + def _build(self, lr_schedule: LearningRateSchedule) -> None: """ Create the network and the optimizer. @@ -210,7 +211,7 @@ class CnnPolicy(DQNPolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[List[int]] = None, activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index 4ea6253..b9ae9fb 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, Optional, Type, Union +from typing import Any, Dict, Optional, Type, Union import numpy as np import torch as th @@ -8,7 +8,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm from stable_baselines3.common.policies import ActorCriticPolicy -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import explained_variance, get_schedule_fn @@ -66,7 +66,7 @@ class PPO(OnPolicyAlgorithm): self, policy: Union[str, Type[ActorCriticPolicy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable] = 3e-4, + learning_rate: Union[float, LearningRateSchedule] = 3e-4, n_steps: int = 2048, batch_size: Optional[int] = 64, n_epochs: int = 10, diff --git a/stable_baselines3/sac/policies.py b/stable_baselines3/sac/policies.py index c79d679..bdc8b92 100644 --- a/stable_baselines3/sac/policies.py +++ b/stable_baselines3/sac/policies.py @@ -14,6 +14,7 @@ from stable_baselines3.common.torch_layers import ( create_mlp, get_actor_critic_arch, ) +from stable_baselines3.common.type_aliases import LearningRateSchedule # CAP the standard deviation of the actor LOG_STD_MAX = 2 @@ -227,7 +228,7 @@ class SACPolicy(BasePolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, @@ -294,7 +295,7 @@ class SACPolicy(BasePolicy): self._build(lr_schedule) - def _build(self, lr_schedule: Callable) -> None: + def _build(self, lr_schedule: LearningRateSchedule) -> None: self.actor = self.make_actor() self.actor.optimizer = self.optimizer_class(self.actor.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs) @@ -397,7 +398,7 @@ class CnnPolicy(SACPolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 785cc9d..d52cf10 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Dict, List, Optional, Tuple, Type, Union import numpy as np import torch as th @@ -7,7 +7,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import polyak_update from stable_baselines3.sac.policies import SACPolicy @@ -74,7 +74,7 @@ class SAC(OffPolicyAlgorithm): self, policy: Union[str, Type[SACPolicy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable] = 3e-4, + learning_rate: Union[float, LearningRateSchedule] = 3e-4, buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 256, diff --git a/stable_baselines3/td3/policies.py b/stable_baselines3/td3/policies.py index 0083a49..ed5669b 100644 --- a/stable_baselines3/td3/policies.py +++ b/stable_baselines3/td3/policies.py @@ -13,6 +13,7 @@ from stable_baselines3.common.torch_layers import ( create_mlp, get_actor_critic_arch, ) +from stable_baselines3.common.type_aliases import LearningRateSchedule class Actor(BasePolicy): @@ -108,7 +109,7 @@ class TD3Policy(BasePolicy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, @@ -163,7 +164,7 @@ class TD3Policy(BasePolicy): self._build(lr_schedule) - def _build(self, lr_schedule: Callable) -> None: + def _build(self, lr_schedule: LearningRateSchedule) -> None: # Create actor and target # the features extractor should not be shared self.actor = self.make_actor(features_extractor=None) @@ -252,7 +253,7 @@ class CnnPolicy(TD3Policy): self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - lr_schedule: Callable, + lr_schedule: LearningRateSchedule, net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index a1186ec..a0e5004 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Dict, List, Optional, Tuple, Type, Union import numpy as np import torch as th @@ -7,7 +7,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.type_aliases import GymEnv, LearningRateSchedule, MaybeCallback from stable_baselines3.common.utils import polyak_update from stable_baselines3.td3.policies import TD3Policy @@ -62,7 +62,7 @@ class TD3(OffPolicyAlgorithm): self, policy: Union[str, Type[TD3Policy]], env: Union[GymEnv, str], - learning_rate: Union[float, Callable] = 1e-3, + learning_rate: Union[float, LearningRateSchedule] = 1e-3, buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 100,