diff --git a/rllib/utils/exploration/epsilon_greedy.py b/rllib/utils/exploration/epsilon_greedy.py index a55849fc6..1aeef99c1 100644 --- a/rllib/utils/exploration/epsilon_greedy.py +++ b/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) diff --git a/rllib/utils/exploration/exploration.py b/rllib/utils/exploration/exploration.py index 0285b699a..a612df235 100644 --- a/rllib/utils/exploration/exploration.py +++ b/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 diff --git a/rllib/utils/exploration/gaussian_noise.py b/rllib/utils/exploration/gaussian_noise.py index 4b47a23cc..549e00ffd 100644 --- a/rllib/utils/exploration/gaussian_noise.py +++ b/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: diff --git a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py b/rllib/utils/exploration/ornstein_uhlenbeck_noise.py index 9cbc0dba9..cde44b715 100644 --- a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py +++ b/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 diff --git a/rllib/utils/exploration/parameter_noise.py b/rllib/utils/exploration/parameter_noise.py index 659f530a0..83d63cf55 100644 --- a/rllib/utils/exploration/parameter_noise.py +++ b/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. diff --git a/rllib/utils/exploration/per_worker_epsilon_greedy.py b/rllib/utils/exploration/per_worker_epsilon_greedy.py index 26d43367a..038fb5b67 100644 --- a/rllib/utils/exploration/per_worker_epsilon_greedy.py +++ b/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. diff --git a/rllib/utils/exploration/per_worker_gaussian_noise.py b/rllib/utils/exploration/per_worker_gaussian_noise.py index d7f72c446..2cfae1230 100644 --- a/rllib/utils/exploration/per_worker_gaussian_noise.py +++ b/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: diff --git a/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py b/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py index 47ec5b155..618fccd79 100644 --- a/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py +++ b/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: diff --git a/rllib/utils/exploration/random.py b/rllib/utils/exploration/random.py index f2f9dee20..1a5549be9 100644 --- a/rllib/utils/exploration/random.py +++ b/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( diff --git a/rllib/utils/exploration/soft_q.py b/rllib/utils/exploration/soft_q.py index aa3086c54..463859427 100644 --- a/rllib/utils/exploration/soft_q.py +++ b/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". """ diff --git a/rllib/utils/exploration/stochastic_sampling.py b/rllib/utils/exploration/stochastic_sampling.py index 410cb9c82..655fe7173 100644 --- a/rllib/utils/exploration/stochastic_sampling.py +++ b/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