[RLlib] Exploration class type annotations. (#11251)

rllib/utils/exploration/epsilon_greedy.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Union, Optional
 
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.utils.annotations import override
@@ -25,10 +25,10 @@ class EpsilonGreedy(Exploration):
                  action_space,
                  *,
                  framework: str,
-                 initial_epsilon=1.0,
-                 final_epsilon=0.05,
-                 epsilon_timesteps=int(1e5),
-                 epsilon_schedule=None,
+                 initial_epsilon: float = 1.0,
+                 final_epsilon: float = 0.05,
+                 epsilon_timesteps: int = int(1e5),
+                 epsilon_schedule: Optional[Schedule] = None,
                  **kwargs):
         """Create an EpsilonGreedy exploration class.
 
@@ -75,7 +75,9 @@ class EpsilonGreedy(Exploration):
             return self._get_torch_exploration_action(q_values, explore,
                                                       timestep)
 
-    def _get_tf_exploration_action_op(self, q_values, explore, timestep):
+    def _get_tf_exploration_action_op(self, q_values: TensorType,
+                                      explore: Union[bool, TensorType],
+                                      timestep: Union[int, TensorType]):
         """TF method to produce the tf op for an epsilon exploration action.
 
         Args:
@@ -119,7 +121,9 @@ class EpsilonGreedy(Exploration):
             with tf1.control_dependencies([assign_op]):
                 return action, tf.zeros_like(action, dtype=tf.float32)
 
-    def _get_torch_exploration_action(self, q_values, explore, timestep):
+    def _get_torch_exploration_action(self, q_values: TensorType,
+                                      explore: bool,
+                                      timestep: Union[int, TensorType]):
         """Torch method to produce an epsilon exploration action.
 
         Args:
@@ -157,7 +161,7 @@ class EpsilonGreedy(Exploration):
             return exploit_action, action_logp
 
     @override(Exploration)
-    def get_info(self, sess=None):
+    def get_info(self, sess: Optional["tf.Session"] = None):
         if sess:
             return sess.run(self._tf_info_op)
         eps = self.epsilon_schedule(self.last_timestep)

rllib/utils/exploration/exploration.py

@@ -1,6 +1,7 @@
 from gym.spaces import Space
 from typing import List, Optional, Union, TYPE_CHECKING
 
+from ray.rllib.env.base_env import BaseEnv
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.policy.sample_batch import SampleBatch
@@ -106,11 +107,11 @@ class Exploration:
 
     @DeveloperAPI
     def on_episode_start(self,
-                         policy,
+                         policy: "Policy",
                          *,
-                         environment=None,
-                         episode=None,
-                         tf_sess=None):
+                         environment: BaseEnv = None,
+                         episode: int = None,
+                         tf_sess: Optional["tf.Session"] = None):
         """Handles necessary exploration logic at the beginning of an episode.
 
         Args:
@@ -123,11 +124,11 @@ class Exploration:
 
     @DeveloperAPI
     def on_episode_end(self,
-                       policy,
+                       policy: "Policy",
                        *,
-                       environment=None,
-                       episode=None,
-                       tf_sess=None):
+                       environment: BaseEnv = None,
+                       episode: int = None,
+                       tf_sess: Optional["tf.Session"] = None):
         """Handles necessary exploration logic at the end of an episode.
 
         Args:
@@ -141,8 +142,8 @@ class Exploration:
     @DeveloperAPI
     def postprocess_trajectory(self,
                                policy: "Policy",
-                               sample_batch,
-                               tf_sess=None):
+                               sample_batch: SampleBatch,
+                               tf_sess: Optional["tf.Session"] = None):
         """Handles post-processing of done episode trajectories.
 
         Changes the given batch in place. This callback is invoked by the
@@ -193,7 +194,7 @@ class Exploration:
         return policy_loss
 
     @DeveloperAPI
-    def get_info(self, sess=None):
+    def get_info(self, sess: Optional["tf.Session"] = None):
         """Returns a description of the current exploration state.
 
         This is not necessarily the state itself (and cannot be used in

rllib/utils/exploration/gaussian_noise.py

@@ -1,4 +1,5 @@
-from typing import Union
+from gym.spaces import Space
+from typing import Union, Optional
 
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.modelv2 import ModelV2
@@ -7,6 +8,7 @@ from ray.rllib.utils.exploration.exploration import Exploration
 from ray.rllib.utils.exploration.random import Random
 from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
     get_variable, TensorType
+from ray.rllib.utils.schedules import Schedule
 from ray.rllib.utils.schedules.piecewise_schedule import PiecewiseSchedule
 
 tf1, tf, tfv = try_import_tf()
@@ -23,16 +25,16 @@ class GaussianNoise(Exploration):
     """
 
     def __init__(self,
-                 action_space,
+                 action_space: Space,
                  *,
                  framework: str,
                  model: ModelV2,
-                 random_timesteps=1000,
-                 stddev=0.1,
-                 initial_scale=1.0,
-                 final_scale=0.02,
-                 scale_timesteps=10000,
-                 scale_schedule=None,
+                 random_timesteps: int = 1000,
+                 stddev: float = 0.1,
+                 initial_scale: float = 1.0,
+                 final_scale: float = 0.02,
+                 scale_timesteps: int = 10000,
+                 scale_schedule: Optional[Schedule] = None,
                  **kwargs):
         """Initializes a GaussianNoise Exploration object.
 
@@ -92,7 +94,9 @@ class GaussianNoise(Exploration):
             return self._get_tf_exploration_action_op(action_distribution,
                                                       explore, timestep)
 
-    def _get_tf_exploration_action_op(self, action_dist, explore, timestep):
+    def _get_tf_exploration_action_op(self, action_dist: ActionDistribution,
+                                      explore: bool,
+                                      timestep: Union[int, TensorType]):
         ts = timestep if timestep is not None else self.last_timestep
 
         # The deterministic actions (if explore=False).
@@ -139,7 +143,9 @@ class GaussianNoise(Exploration):
             with tf1.control_dependencies([assign_op]):
                 return action, logp
 
-    def _get_torch_exploration_action(self, action_dist, explore, timestep):
+    def _get_torch_exploration_action(self, action_dist: ActionDistribution,
+                                      explore: bool,
+                                      timestep: Union[int, TensorType]):
         # Set last timestep or (if not given) increase by one.
         self.last_timestep = timestep if timestep is not None else \
             self.last_timestep + 1
@@ -180,7 +186,7 @@ class GaussianNoise(Exploration):
         return action, logp
 
     @override(Exploration)
-    def get_info(self, sess=None):
+    def get_info(self, sess: Optional["tf.Session"] = None):
         """Returns the current scale value.
 
         Returns:

rllib/utils/exploration/ornstein_uhlenbeck_noise.py

@@ -1,9 +1,12 @@
 import numpy as np
+from typing import Optional, Union
 
+from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise
 from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
-    get_variable
+    get_variable, TensorType
+from ray.rllib.utils.schedules import Schedule
 
 tf1, tf, tfv = try_import_tf()
 torch, _ = try_import_torch()
@@ -24,14 +27,14 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
                  action_space,
                  *,
                  framework: str,
-                 ou_theta=0.15,
-                 ou_sigma=0.2,
-                 ou_base_scale=0.1,
-                 random_timesteps=1000,
-                 initial_scale=1.0,
-                 final_scale=0.02,
-                 scale_timesteps=10000,
-                 scale_schedule=None,
+                 ou_theta: float = 0.15,
+                 ou_sigma: float = 0.2,
+                 ou_base_scale: float = 0.1,
+                 random_timesteps: int = 1000,
+                 initial_scale: float = 1.0,
+                 final_scale: float = 0.02,
+                 scale_timesteps: int = 10000,
+                 scale_schedule: Optional[Schedule] = None,
                  **kwargs):
         """Initializes an Ornstein-Uhlenbeck Exploration object.
 
@@ -82,7 +85,9 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
             device=self.device)
 
     @override(GaussianNoise)
-    def _get_tf_exploration_action_op(self, action_dist, explore, timestep):
+    def _get_tf_exploration_action_op(self, action_dist: ActionDistribution,
+                                      explore: Union[bool, TensorType],
+                                      timestep: Union[int, TensorType]):
         ts = timestep if timestep is not None else self.last_timestep
         scale = self.scale_schedule(ts)
 
@@ -143,7 +148,9 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
                 return action, logp
 
     @override(GaussianNoise)
-    def _get_torch_exploration_action(self, action_dist, explore, timestep):
+    def _get_torch_exploration_action(self, action_dist: ActionDistribution,
+                                      explore: bool,
+                                      timestep: Union[int, TensorType]):
         # Set last timestep or (if not given) increase by one.
         self.last_timestep = timestep if timestep is not None else \
             self.last_timestep + 1

rllib/utils/exploration/parameter_noise.py

@@ -1,17 +1,24 @@
 from gym.spaces import Box, Discrete
 import numpy as np
+from typing import Optional, TYPE_CHECKING, Union
 
-from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.env.base_env import BaseEnv
+from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.tf.tf_action_dist import Categorical, Deterministic
 from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
     TorchDeterministic
+from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.exploration import Exploration
-from ray.rllib.utils.framework import try_import_tf, try_import_torch
-from ray.rllib.utils.framework import get_variable
+from ray.rllib.utils.framework import get_variable, try_import_tf, \
+    try_import_torch
 from ray.rllib.utils.from_config import from_config
 from ray.rllib.utils.numpy import softmax, SMALL_NUMBER
+from ray.rllib.utils.typing import TensorType
+
+if TYPE_CHECKING:
+    from ray.rllib.policy.policy import Policy
 
 tf1, tf, tfv = try_import_tf()
 torch, _ = try_import_torch()
@@ -36,9 +43,9 @@ class ParameterNoise(Exploration):
                  framework: str,
                  policy_config: dict,
                  model: ModelV2,
-                 initial_stddev=1.0,
-                 random_timesteps=10000,
-                 sub_exploration=None,
+                 initial_stddev: float = 1.0,
+                 random_timesteps: int = 10000,
+                 sub_exploration: Optional[dict] = None,
                  **kwargs):
         """Initializes a ParameterNoise Exploration object.
 
@@ -139,9 +146,9 @@ class ParameterNoise(Exploration):
     @override(Exploration)
     def before_compute_actions(self,
                                *,
-                               timestep=None,
-                               explore=None,
-                               tf_sess=None):
+                               timestep: Optional[int] = None,
+                               explore: bool = None,
+                               tf_sess: Optional["tf.Session"] = None):
         explore = explore if explore is not None else \
             self.policy_config["explore"]
 
@@ -158,11 +165,10 @@ class ParameterNoise(Exploration):
             self._remove_noise(tf_sess=tf_sess)
 
     @override(Exploration)
-    def get_exploration_action(self,
-                               *,
-                               action_distribution,
-                               timestep,
-                               explore=True):
+    def get_exploration_action(self, *,
+                               action_distribution: ActionDistribution,
+                               timestep: Union[TensorType, int],
+                               explore: Union[TensorType, bool]):
         # Use our sub-exploration object to handle the final exploration
         # action (depends on the algo-type/action-space/etc..).
         return self.sub_exploration.get_exploration_action(
@@ -172,11 +178,11 @@ class ParameterNoise(Exploration):
 
     @override(Exploration)
     def on_episode_start(self,
-                         policy,
+                         policy: "Policy",
                          *,
-                         environment=None,
-                         episode=None,
-                         tf_sess=None):
+                         environment: BaseEnv = None,
+                         episode: int = None,
+                         tf_sess: Optional["tf.Session"] = None):
         # We have to delay the noise-adding step by one forward call.
         # This is due to the fact that the optimizer does it's step right
         # after the episode was reset (and hence the noise was already added!).
@@ -204,7 +210,10 @@ class ParameterNoise(Exploration):
             self._remove_noise(tf_sess=tf_sess)
 
     @override(Exploration)
-    def postprocess_trajectory(self, policy, sample_batch, tf_sess=None):
+    def postprocess_trajectory(self,
+                               policy: "Policy",
+                               sample_batch: SampleBatch,
+                               tf_sess: Optional["tf.Session"] = None):
         noisy_action_dist = noise_free_action_dist = None
         # Adjust the stddev depending on the action (pi)-distance.
         # Also see [1] for details.

rllib/utils/exploration/per_worker_epsilon_greedy.py

@@ -1,3 +1,6 @@
+from gym.spaces import Space
+from typing import Optional
+
 from ray.rllib.utils.exploration.epsilon_greedy import EpsilonGreedy
 from ray.rllib.utils.schedules import ConstantSchedule
 
@@ -10,7 +13,8 @@ class PerWorkerEpsilonGreedy(EpsilonGreedy):
     See Ape-X paper.
     """
 
-    def __init__(self, action_space, *, framework, num_workers, worker_index,
+    def __init__(self, action_space: Space, *, framework: str,
+                 num_workers: Optional[int], worker_index: Optional[int],
                  **kwargs):
         """Create a PerWorkerEpsilonGreedy exploration class.
 

rllib/utils/exploration/per_worker_gaussian_noise.py

@@ -1,3 +1,6 @@
+from gym.spaces import Space
+from typing import Optional
+
 from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise
 from ray.rllib.utils.schedules import ConstantSchedule
 
@@ -10,7 +13,8 @@ class PerWorkerGaussianNoise(GaussianNoise):
     See Ape-X paper.
     """
 
-    def __init__(self, action_space, *, framework, num_workers, worker_index,
+    def __init__(self, action_space: Space, *, framework: Optional[str],
+                 num_workers: Optional[int], worker_index: Optional[int],
                  **kwargs):
         """
         Args:

rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py

@@ -1,3 +1,6 @@
+from gym.spaces import Space
+from typing import Optional
+
 from ray.rllib.utils.exploration.ornstein_uhlenbeck_noise import \
     OrnsteinUhlenbeckNoise
 from ray.rllib.utils.schedules import ConstantSchedule
@@ -11,7 +14,8 @@ class PerWorkerOrnsteinUhlenbeckNoise(OrnsteinUhlenbeckNoise):
     See Ape-X paper.
     """
 
-    def __init__(self, action_space, *, framework, num_workers, worker_index,
+    def __init__(self, action_space: Space, *, framework: Optional[str],
+                 num_workers: Optional[int], worker_index: Optional[int],
                  **kwargs):
         """
         Args:

rllib/utils/exploration/random.py

@@ -1,9 +1,10 @@
-from gym.spaces import Discrete, Box, MultiDiscrete
+from gym.spaces import Discrete, Box, MultiDiscrete, Space
 import numpy as np
 import tree
-from typing import Union
+from typing import Union, Optional
 
 from ray.rllib.models.action_dist import ActionDistribution
+from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.exploration import Exploration
 from ray.rllib.utils import force_tuple
@@ -23,7 +24,8 @@ class Random(Exploration):
     If explore=False, returns the greedy/max-likelihood action.
     """
 
-    def __init__(self, action_space, *, model, framework, **kwargs):
+    def __init__(self, action_space: Space, *, model: ModelV2,
+                 framework: Optional[str], **kwargs):
         """Initialize a Random Exploration object.
 
         Args:
@@ -53,7 +55,9 @@ class Random(Exploration):
             return self.get_torch_exploration_action(action_distribution,
                                                      explore)
 
-    def get_tf_exploration_action_op(self, action_dist, explore):
+    def get_tf_exploration_action_op(
+            self, action_dist: ActionDistribution,
+            explore: Optional[Union[bool, TensorType]]):
         def true_fn():
             batch_size = 1
             req = force_tuple(
@@ -111,7 +115,8 @@ class Random(Exploration):
         logp = tf.zeros(shape=(batch_size, ), dtype=tf.float32)
         return action, logp
 
-    def get_torch_exploration_action(self, action_dist, explore):
+    def get_torch_exploration_action(self, action_dist: ActionDistribution,
+                                     explore: bool):
         if explore:
             req = force_tuple(
                 action_dist.required_model_output_shape(

rllib/utils/exploration/soft_q.py

@@ -1,5 +1,5 @@
-from gym.spaces import Discrete
-from typing import Union
+from gym.spaces import Discrete, Space
+from typing import Union, Optional
 
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.tf.tf_action_dist import Categorical
@@ -16,12 +16,17 @@ class SoftQ(StochasticSampling):
     output divided by the temperature. Returns the argmax iff explore=False.
     """
 
-    def __init__(self, action_space, *, framework, temperature=1.0, **kwargs):
+    def __init__(self,
+                 action_space: Space,
+                 *,
+                 framework: Optional[str],
+                 temperature: float = 1.0,
+                 **kwargs):
         """Initializes a SoftQ Exploration object.
 
         Args:
             action_space (Space): The gym action space used by the environment.
-            temperature (Schedule): The temperature to divide model outputs by
+            temperature (float): The temperature to divide model outputs by
                 before creating the Categorical distribution to sample from.
             framework (str): One of None, "tf", "torch".
         """

rllib/utils/exploration/stochastic_sampling.py

@@ -1,3 +1,4 @@
+import gym
 import tree
 from typing import Union
 
@@ -23,7 +24,7 @@ class StochasticSampling(Exploration):
     """
 
     def __init__(self,
-                 action_space,
+                 action_space: gym.spaces.Space,
                  *,
                  framework: str,
                  model: ModelV2,
@@ -32,7 +33,8 @@ class StochasticSampling(Exploration):
         """Initializes a StochasticSampling Exploration object.
 
         Args:
-            action_space (Space): The gym action space used by the environment.
+            action_space (gym.spaces.Space): The gym action space used by the
+                environment.
             framework (str): One of None, "tf", "torch".
             model (ModelV2): The ModelV2 used by the owning Policy.
             random_timesteps (int): The number of timesteps for which to act
@@ -108,7 +110,9 @@ class StochasticSampling(Exploration):
             with tf1.control_dependencies([assign_op]):
                 return action, logp
 
-    def _get_torch_exploration_action(self, action_dist, timestep, explore):
+    def _get_torch_exploration_action(self, action_dist: ActionDistribution,
+                                      timestep: Union[TensorType, int],
+                                      explore: Union[TensorType, bool]):
         # Set last timestep or (if not given) increase by one.
         self.last_timestep = timestep if timestep is not None else \
             self.last_timestep + 1