diff --git a/rllib/agents/a3c/a3c_torch_policy.py b/rllib/agents/a3c/a3c_torch_policy.py index 36320d3bb..f12e817e0 100644 --- a/rllib/agents/a3c/a3c_torch_policy.py +++ b/rllib/agents/a3c/a3c_torch_policy.py @@ -2,15 +2,15 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import torch -import torch.nn.functional as F -from torch import nn - import ray from ray.rllib.evaluation.postprocessing import compute_advantages, \ Postprocessing from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.torch_policy_template import build_torch_policy +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() +F = nn.functional def actor_critic_loss(policy, model, dist_class, train_batch): diff --git a/rllib/agents/qmix/mixers.py b/rllib/agents/qmix/mixers.py index 3f8fbbce4..1e91622c0 100644 --- a/rllib/agents/qmix/mixers.py +++ b/rllib/agents/qmix/mixers.py @@ -2,18 +2,20 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import torch as th -import torch.nn as nn -import torch.nn.functional as F import numpy as np +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() +F = nn.functional + class VDNMixer(nn.Module): def __init__(self): super(VDNMixer, self).__init__() def forward(self, agent_qs, batch): - return th.sum(agent_qs, dim=2, keepdim=True) + return torch.sum(agent_qs, dim=2, keepdim=True) class QMixer(nn.Module): @@ -47,18 +49,18 @@ class QMixer(nn.Module): states = states.reshape(-1, self.state_dim) agent_qs = agent_qs.view(-1, 1, self.n_agents) # First layer - w1 = th.abs(self.hyper_w_1(states)) + w1 = torch.abs(self.hyper_w_1(states)) b1 = self.hyper_b_1(states) w1 = w1.view(-1, self.n_agents, self.embed_dim) b1 = b1.view(-1, 1, self.embed_dim) - hidden = F.elu(th.bmm(agent_qs, w1) + b1) + hidden = F.elu(torch.bmm(agent_qs, w1) + b1) # Second layer - w_final = th.abs(self.hyper_w_final(states)) + w_final = torch.abs(self.hyper_w_final(states)) w_final = w_final.view(-1, self.embed_dim, 1) # State-dependent bias v = self.V(states).view(-1, 1, 1) # Compute final output - y = th.bmm(hidden, w_final) + v + y = torch.bmm(hidden, w_final) + v # Reshape and return q_tot = y.view(bs, -1, 1) return q_tot diff --git a/rllib/agents/qmix/model.py b/rllib/agents/qmix/model.py index 059d7fa2e..e5e5e9c86 100644 --- a/rllib/agents/qmix/model.py +++ b/rllib/agents/qmix/model.py @@ -2,12 +2,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from torch import nn -import torch.nn.functional as F - from ray.rllib.models.preprocessors import get_preprocessor from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.utils.annotations import override +from ray.rllib.utils import try_import_torch + +torch, nn = try_import_torch() +F = nn.functional class RNNModel(TorchModelV2, nn.Module): diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 06ab3d35a..7f068c4af 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -158,7 +158,7 @@ COMMON_CONFIG = { # after the initial eager pass. "eager_tracing": False, # Disable eager execution on workers (but allow it on the driver). This - # only has an effect is eager is enabled. + # only has an effect if eager is enabled. "no_eager_on_workers": False, # === Evaluation Settings === diff --git a/rllib/agents/trainer_template.py b/rllib/agents/trainer_template.py index dae1de45e..638e92ec4 100644 --- a/rllib/agents/trainer_template.py +++ b/rllib/agents/trainer_template.py @@ -8,9 +8,6 @@ from ray.rllib.agents.trainer import Trainer, COMMON_CONFIG from ray.rllib.optimizers import SyncSamplesOptimizer from ray.rllib.utils import add_mixins from ray.rllib.utils.annotations import override, DeveloperAPI -from ray.rllib.utils import try_import_tf - -tf = try_import_tf() @DeveloperAPI @@ -161,7 +158,6 @@ def build_trainer(name, Trainer.__setstate__(self, state) self.state = state["trainer_state"].copy() - @staticmethod def with_updates(**overrides): """Build a copy of this trainer with the specified overrides. @@ -171,7 +167,7 @@ def build_trainer(name, """ return build_trainer(**dict(original_kwargs, **overrides)) - trainer_cls.with_updates = with_updates + trainer_cls.with_updates = staticmethod(with_updates) trainer_cls.__name__ = name trainer_cls.__qualname__ = name return trainer_cls diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py index 2b1a29059..4ed7adbb6 100644 --- a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py +++ b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py @@ -1,10 +1,12 @@ import numpy as np -import torch + from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.torch_policy import TorchPolicy from ray.rllib.utils.annotations import override - from ray.rllib.contrib.alpha_zero.core.mcts import Node, RootParentNode +from ray.rllib.utils import try_import_torch + +torch, _ = try_import_torch() class AlphaZeroPolicy(TorchPolicy): diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py b/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py index 0b3d9a9e9..5c19dcdf4 100644 --- a/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py +++ b/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py @@ -4,15 +4,13 @@ from __future__ import print_function import logging -import torch -import torch.nn as nn from ray.rllib.agents import with_common_config from ray.rllib.agents.trainer_template import build_trainer from ray.rllib.models.catalog import ModelCatalog from ray.rllib.models.model import restore_original_dimensions from ray.rllib.models.torch.torch_action_dist import TorchCategorical from ray.rllib.optimizers import SyncSamplesOptimizer -from ray.rllib.utils import try_import_tf +from ray.rllib.utils import try_import_tf, try_import_torch from ray.tune.registry import ENV_CREATOR, _global_registry from ray.rllib.contrib.alpha_zero.core.alpha_zero_policy import AlphaZeroPolicy @@ -22,6 +20,8 @@ from ray.rllib.contrib.alpha_zero.optimizer.sync_batches_replay_optimizer \ import SyncBatchesReplayOptimizer tf = try_import_tf() +torch, nn = try_import_torch() + logger = logging.getLogger(__name__) diff --git a/rllib/contrib/alpha_zero/models/custom_torch_models.py b/rllib/contrib/alpha_zero/models/custom_torch_models.py index bf0291ee4..cc4cec7da 100644 --- a/rllib/contrib/alpha_zero/models/custom_torch_models.py +++ b/rllib/contrib/alpha_zero/models/custom_torch_models.py @@ -1,11 +1,13 @@ from abc import ABC import numpy as np -import torch -import torch.nn as nn + from ray.rllib.models.model import restore_original_dimensions from ray.rllib.models.preprocessors import get_preprocessor from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils import try_import_torch + +torch, nn = try_import_torch() def convert_to_tensor(arr): diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index 90d5e5bcb..f40521ddd 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -32,9 +32,11 @@ from ray.rllib.utils.debug import disable_log_once_globally, log_once, \ summarize, enable_periodic_logging from ray.rllib.utils.filter import get_filter from ray.rllib.utils.tf_run_builder import TFRunBuilder -from ray.rllib.utils import try_import_tf +from ray.rllib.utils import try_import_tf, try_import_torch tf = try_import_tf() +torch, _ = try_import_torch() + logger = logging.getLogger(__name__) # Handle to the current rollout worker, which will be set to the most recently @@ -321,9 +323,9 @@ class RolloutWorker(EvaluatorInterface): self.env)) self.env.seed(seed) try: - import torch + assert torch is not None torch.manual_seed(seed) - except ImportError: + except AssertionError: logger.info("Could not seed torch") if _has_tensorflow_graph(policy_dict) and not (tf and tf.executing_eagerly()): diff --git a/rllib/models/model.py b/rllib/models/model.py index 242fb2ea2..996179630 100644 --- a/rllib/models/model.py +++ b/rllib/models/model.py @@ -9,9 +9,11 @@ import gym from ray.rllib.models.tf.misc import linear, normc_initializer from ray.rllib.models.preprocessors import get_preprocessor from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI -from ray.rllib.utils import try_import_tf +from ray.rllib.utils import try_import_tf, try_import_torch tf = try_import_tf() +torch, _ = try_import_torch() + logger = logging.getLogger(__name__) @@ -196,7 +198,7 @@ def flatten(obs, framework): if framework == "tf": return tf.layers.flatten(obs) elif framework == "torch": - import torch + assert torch is not None return torch.flatten(obs, start_dim=1) else: raise NotImplementedError("flatten", framework) @@ -225,7 +227,7 @@ def restore_original_dimensions(obs, obs_space, tensorlib=tf): if tensorlib == "tf": tensorlib = tf elif tensorlib == "torch": - import torch + assert torch is not None tensorlib = torch return _unpack_obs(obs, obs_space.original_space, tensorlib=tensorlib) else: diff --git a/rllib/models/tf/misc.py b/rllib/models/tf/misc.py index 73ee1d87c..0e62783f3 100644 --- a/rllib/models/tf/misc.py +++ b/rllib/models/tf/misc.py @@ -18,6 +18,8 @@ def normc_initializer(std=1.0): def get_activation_fn(name): + if name == "linear": + return None return getattr(tf.nn, name) diff --git a/rllib/models/torch/fcnet.py b/rllib/models/torch/fcnet.py index 6b1e487f0..b168f2ad2 100644 --- a/rllib/models/torch/fcnet.py +++ b/rllib/models/torch/fcnet.py @@ -4,12 +4,14 @@ from __future__ import print_function import logging import numpy as np -import torch.nn as nn from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.misc import normc_initializer, SlimFC, \ _get_activation_fn from ray.rllib.utils.annotations import override +from ray.rllib.utils import try_import_torch + +_, nn = try_import_torch() logger = logging.getLogger(__name__) diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index e9051d9f9..806c646dd 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -4,8 +4,10 @@ from __future__ import division from __future__ import print_function import numpy as np -import torch -import torch.nn as nn + +from ray.rllib.utils import try_import_torch + +torch, nn = try_import_torch() def normc_initializer(std=1.0): @@ -51,14 +53,14 @@ def valid_padding(in_size, filter_size, stride_size): def _get_activation_fn(name): - activation = None if name == "tanh": - activation = nn.Tanh + return nn.Tanh elif name == "relu": - activation = nn.ReLU + return nn.ReLU + elif name == "linear": + return None else: raise ValueError("Unknown activation: {}".format(name)) - return activation class SlimConv2d(nn.Module): diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index 6c4eb7268..af0962733 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/rllib/models/torch/torch_action_dist.py @@ -2,15 +2,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -try: - import torch -except ImportError: - pass # soft dep - import numpy as np from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.utils.annotations import override +from ray.rllib.utils import try_import_torch + +torch, nn = try_import_torch() class TorchDistributionWrapper(ActionDistribution): @@ -26,7 +24,7 @@ class TorchDistributionWrapper(ActionDistribution): @override(ActionDistribution) def kl(self, other): - return torch.distributions.kl.kl_divergence(self.dist, other) + return torch.distributions.kl.kl_divergence(self.dist, other.dist) @override(ActionDistribution) def sample(self): diff --git a/rllib/models/torch/torch_modelv2.py b/rllib/models/torch/torch_modelv2.py index 933e59917..cd279c7fc 100644 --- a/rllib/models/torch/torch_modelv2.py +++ b/rllib/models/torch/torch_modelv2.py @@ -2,10 +2,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import torch.nn as nn - from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils import try_import_torch + +_, nn = try_import_torch() @PublicAPI diff --git a/rllib/models/torch/visionnet.py b/rllib/models/torch/visionnet.py index c0e74dc51..077809802 100644 --- a/rllib/models/torch/visionnet.py +++ b/rllib/models/torch/visionnet.py @@ -2,13 +2,14 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import torch.nn as nn - from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.misc import normc_initializer, valid_padding, \ SlimConv2d, SlimFC from ray.rllib.models.tf.visionnet_v1 import _get_filter_config from ray.rllib.utils.annotations import override +from ray.rllib.utils import try_import_torch + +_, nn = try_import_torch() class VisionNetwork(TorchModelV2, nn.Module): diff --git a/rllib/tests/test_nested_spaces.py b/rllib/tests/test_nested_spaces.py index 793d24710..d3283d54d 100644 --- a/rllib/tests/test_nested_spaces.py +++ b/rllib/tests/test_nested_spaces.py @@ -7,7 +7,6 @@ import pickle from gym import spaces from gym.envs.registration import EnvSpec import gym -import torch.nn as nn import unittest import ray @@ -24,9 +23,11 @@ from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.rollout import rollout from ray.rllib.tests.test_external_env import SimpleServing from ray.tune.registry import register_env -from ray.rllib.utils import try_import_tf +from ray.rllib.utils import try_import_tf, try_import_torch tf = try_import_tf() +_, nn = try_import_torch() + DICT_SPACE = spaces.Dict({ "sensors": spaces.Dict({ diff --git a/rllib/utils/__init__.py b/rllib/utils/__init__.py index c82e324d5..bdfb73b25 100644 --- a/rllib/utils/__init__.py +++ b/rllib/utils/__init__.py @@ -1,31 +1,17 @@ -import logging -import os - +from ray.rllib.utils.annotations import override, PublicAPI, DeveloperAPI +from ray.rllib.utils.framework import try_import_tf, try_import_tfp, \ + try_import_torch +from ray.rllib.utils.deprecation import deprecation_warning, renamed_agent, \ + renamed_class, renamed_function from ray.rllib.utils.filter_manager import FilterManager from ray.rllib.utils.filter import Filter +from ray.rllib.utils.numpy import sigmoid, softmax, relu, one_hot, fc, lstm, \ + SMALL_NUMBER, LARGE_INTEGER from ray.rllib.utils.policy_client import PolicyClient from ray.rllib.utils.policy_server import PolicyServer +from ray.rllib.utils.test_utils import check from ray.tune.util import merge_dicts, deep_update -logger = logging.getLogger(__name__) - - -def renamed_class(cls, old_name): - """Helper class for renaming classes with a warning.""" - - class DeprecationWrapper(cls): - # note: **kw not supported for ray.remote classes - def __init__(self, *args, **kw): - new_name = cls.__module__ + "." + cls.__name__ - logger.warning("DeprecationWarning: {} has been renamed to {}. ". - format(old_name, new_name) + - "This will raise an error in the future.") - cls.__init__(self, *args, **kw) - - DeprecationWrapper.__name__ = cls.__name__ - - return DeprecationWrapper - def add_mixins(base, mixins): """Returns a new class with mixins applied in priority order.""" @@ -42,63 +28,31 @@ def add_mixins(base, mixins): return base -def renamed_agent(cls): - """Helper class for renaming Agent => Trainer with a warning.""" - - class DeprecationWrapper(cls): - def __init__(self, config=None, env=None, logger_creator=None): - old_name = cls.__name__.replace("Trainer", "Agent") - new_name = cls.__module__ + "." + cls.__name__ - logger.warning("DeprecationWarning: {} has been renamed to {}. ". - format(old_name, new_name) + - "This will raise an error in the future.") - cls.__init__(self, config, env, logger_creator) - - DeprecationWrapper.__name__ = cls.__name__ - - return DeprecationWrapper - - -def try_import_tf(): - if "RLLIB_TEST_NO_TF_IMPORT" in os.environ: - logger.warning("Not importing TensorFlow for test purposes") - return None - - try: - if "TF_CPP_MIN_LOG_LEVEL" not in os.environ: - os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" - import tensorflow.compat.v1 as tf - tf.logging.set_verbosity(tf.logging.ERROR) - tf.disable_v2_behavior() - return tf - except ImportError: - try: - import tensorflow as tf - return tf - except ImportError: - return None - - -def try_import_tfp(): - if "RLLIB_TEST_NO_TF_IMPORT" in os.environ: - logger.warning( - "Not importing TensorFlow Probability for test purposes.") - return None - - try: - import tensorflow_probability as tfp - return tfp - except ImportError: - return None - - __all__ = [ - "Filter", - "FilterManager", - "PolicyClient", - "PolicyServer", - "merge_dicts", + "add_mixins", + "check", + "deprecation_warning", + "fc", + "lstm", + "one_hot", + "relu", + "sigmoid", + "softmax", "deep_update", + "merge_dicts", + "override", + "renamed_function", + "renamed_agent", "renamed_class", "try_import_tf", + "try_import_tfp", + "try_import_torch", + "DeveloperAPI", + "Filter", + "FilterManager", + "LARGE_INTEGER", + "PolicyClient", + "PolicyServer", + "PublicAPI", + "SMALL_NUMBER", ] diff --git a/rllib/utils/deprecation.py b/rllib/utils/deprecation.py new file mode 100644 index 000000000..707eebcdb --- /dev/null +++ b/rllib/utils/deprecation.py @@ -0,0 +1,66 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import logging + + +logger = logging.getLogger(__name__) + + +def deprecation_warning(old, new=None): + logger.warning( + "DeprecationWarning: `{}` has been deprecated.".format(old) + + (" Use `{}` instead." if new else "") + + " This will raise an error in the future!" + ) + + +def renamed_class(cls, old_name): + """Helper class for renaming classes with a warning.""" + + class DeprecationWrapper(cls): + # note: **kw not supported for ray.remote classes + def __init__(self, *args, **kw): + new_name = cls.__module__ + "." + cls.__name__ + deprecation_warning(old_name, new_name) + cls.__init__(self, *args, **kw) + + DeprecationWrapper.__name__ = cls.__name__ + + return DeprecationWrapper + + +def renamed_agent(cls): + """Helper class for renaming Agent => Trainer with a warning.""" + + class DeprecationWrapper(cls): + def __init__(self, config=None, env=None, logger_creator=None): + old_name = cls.__name__.replace("Trainer", "Agent") + new_name = cls.__module__ + "." + cls.__name__ + deprecation_warning(old_name, new_name) + cls.__init__(self, config, env, logger_creator) + + DeprecationWrapper.__name__ = cls.__name__ + + return DeprecationWrapper + + +def renamed_function(func, old_name): + """Helper function for renaming a function.""" + + def deprecation_wrapper(*args, **kwargs): + new_name = func.__module__ + "." + func.__name__ + deprecation_warning(old_name, new_name) + return func(*args, **kwargs) + + deprecation_wrapper.__name__ = func.__name__ + + return deprecation_wrapper + + +def moved_function(func): + new_location = func.__module__ + deprecation_warning("import {}".format(func.__name__), "import {}". + format(new_location)) + return func diff --git a/rllib/utils/explained_variance.py b/rllib/utils/explained_variance.py index a3e9cbadb..907a00b6b 100644 --- a/rllib/utils/explained_variance.py +++ b/rllib/utils/explained_variance.py @@ -2,12 +2,18 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from ray.rllib.utils import try_import_tf +from ray.rllib.utils import try_import_tf, try_import_torch tf = try_import_tf() +torch, nn = try_import_torch() -def explained_variance(y, pred): - _, y_var = tf.nn.moments(y, axes=[0]) - _, diff_var = tf.nn.moments(y - pred, axes=[0]) - return tf.maximum(-1.0, 1 - (diff_var / y_var)) +def explained_variance(y, pred, framework="tf"): + if framework == "tf": + _, y_var = tf.nn.moments(y, axes=[0]) + _, diff_var = tf.nn.moments(y - pred, axes=[0]) + return tf.maximum(-1.0, 1 - (diff_var / y_var)) + else: + y_var = torch.var(y, dim=[0]) + diff_var = torch.var(y - pred, dim=[0]) + return max(-1.0, 1 - (diff_var / y_var)) diff --git a/rllib/utils/framework.py b/rllib/utils/framework.py new file mode 100644 index 000000000..23cab9fdf --- /dev/null +++ b/rllib/utils/framework.py @@ -0,0 +1,66 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import logging +import os + +logger = logging.getLogger(__name__) + + +def try_import_tf(): + """ + Returns: + The tf module (either from tf2.0.compat.v1 OR as tf1.x. + """ + if "RLLIB_TEST_NO_TF_IMPORT" in os.environ: + logger.warning("Not importing TensorFlow for test purposes") + return None + + try: + if "TF_CPP_MIN_LOG_LEVEL" not in os.environ: + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" + import tensorflow.compat.v1 as tf + tf.logging.set_verbosity(tf.logging.ERROR) + tf.disable_v2_behavior() + return tf + except ImportError: + try: + import tensorflow as tf + return tf + except ImportError: + return None + + +def try_import_tfp(): + """ + Returns: + The tfp module. + """ + if "RLLIB_TEST_NO_TF_IMPORT" in os.environ: + logger.warning("Not importing TensorFlow Probability for test " + "purposes.") + return None + + try: + import tensorflow_probability as tfp + return tfp + except ImportError: + return None + + +def try_import_torch(): + """ + Returns: + tuple: torch AND torch.nn modules. + """ + if "RLLIB_TEST_NO_TORCH_IMPORT" in os.environ: + logger.warning("Not importing Torch for test purposes.") + return None, None + + try: + import torch + import torch.nn as nn + return torch, nn + except ImportError: + return None, None diff --git a/rllib/utils/numpy.py b/rllib/utils/numpy.py new file mode 100644 index 000000000..7b0040dbc --- /dev/null +++ b/rllib/utils/numpy.py @@ -0,0 +1,198 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + + +SMALL_NUMBER = 1e-6 +# Some large int number. May be increased here, if needed. +LARGE_INTEGER = 100000000 +# Min and Max outputs (clipped) from an NN-output layer interpreted as the +# log(x) of some x (e.g. a stddev of a normal +# distribution). +MIN_LOG_NN_OUTPUT = -20 +MAX_LOG_NN_OUTPUT = 2 + + +def sigmoid(x, derivative=False): + """ + Returns the sigmoid function applied to x. + Alternatively, can return the derivative or the sigmoid function. + + Args: + x (np.ndarray): The input to the sigmoid function. + derivative (bool): Whether to return the derivative or not. + Default: False. + + Returns: + np.ndarray: The sigmoid function (or its derivative) applied to x. + """ + if derivative: + return x * (1 - x) + else: + return 1 / (1 + np.exp(-x)) + + +def softmax(x, axis=-1): + """ + Returns the softmax values for x as: + S(xi) = e^xi / SUMj(e^xj), where j goes over all elements in x. + + Args: + x (np.ndarray): The input to the softmax function. + axis (int): The axis along which to softmax. + + Returns: + np.ndarray: The softmax over x. + """ + x_exp = np.exp(x) + return np.maximum(x_exp / np.sum(x_exp, axis, keepdims=True), SMALL_NUMBER) + + +def relu(x, alpha=0.0): + """ + Implementation of the leaky ReLU function: + y = x * alpha if x < 0 else x + + Args: + x (np.ndarray): The input values. + alpha (float): A scaling ("leak") factor to use for negative x. + + Returns: + np.ndarray: The leaky ReLU output for x. + """ + return np.maximum(x, x*alpha, x) + + +def one_hot(x, depth=0, on_value=1, off_value=0): + """ + One-hot utility function for numpy. + Thanks to qianyizhang: + https://gist.github.com/qianyizhang/07ee1c15cad08afb03f5de69349efc30. + + Args: + x (np.ndarray): The input to be one-hot encoded. + depth (int): The max. number to be one-hot encoded (size of last rank). + on_value (float): The value to use for on. Default: 1.0. + off_value (float): The value to use for off. Default: 0.0. + + Returns: + np.ndarray: The one-hot encoded equivalent of the input array. + """ + # Handle bool arrays correctly. + if x.dtype == np.bool_: + x = x.astype(np.int) + depth = 2 + + if depth == 0: + depth = np.max(x) + 1 + assert np.max(x) < depth, \ + "ERROR: The max. index of `x` ({}) is larger than depth ({})!".\ + format(np.max(x), depth) + shape = x.shape + + # Python 2.7 compatibility, (*shape, depth) is not allowed. + shape_list = shape[:] + shape_list.append(depth) + out = np.ones(shape_list) * off_value + indices = [] + for i in range(x.ndim): + tiles = [1] * x.ndim + s = [1] * x.ndim + s[i] = -1 + r = np.arange(shape[i]).reshape(s) + if i > 0: + tiles[i-1] = shape[i-1] + r = np.tile(r, tiles) + indices.append(r) + indices.append(x) + out[tuple(indices)] = on_value + return out + + +def fc(x, weights, biases=None): + """ + Calculates the outputs of a fully-connected (dense) layer given + weights/biases and an input. + + Args: + x (np.ndarray): The input to the dense layer. + weights (np.ndarray): The weights matrix. + biases (Optional[np.ndarray]): The biases vector. All 0s if None. + + Returns: + The dense layer's output. + """ + return np.matmul(x, weights) + (0.0 if biases is None else biases) + + +def lstm(x, weights, biases=None, initial_internal_states=None, + time_major=False, forget_bias=1.0): + """ + Calculates the outputs of an LSTM layer given weights/biases, + internal_states, and input. + + Args: + x (np.ndarray): The inputs to the LSTM layer including time-rank + (0th if time-major, else 1st) and the batch-rank + (1st if time-major, else 0th). + + weights (np.ndarray): The weights matrix. + biases (Optional[np.ndarray]): The biases vector. All 0s if None. + + initial_internal_states (Optional[np.ndarray]): The initial internal + states to pass into the layer. All 0s if None. + + time_major (bool): Whether to use time-major or not. Default: False. + + forget_bias (float): Gets added to first sigmoid (forget gate) output. + Default: 1.0. + + Returns: + Tuple: + - The LSTM layer's output. + - Tuple: Last (c-state, h-state). + """ + sequence_length = x.shape[0 if time_major else 1] + batch_size = x.shape[1 if time_major else 0] + units = weights.shape[1] // 4 # 4 internal layers (3x sigmoid, 1x tanh) + + if initial_internal_states is None: + c_states = np.zeros(shape=(batch_size, units)) + h_states = np.zeros(shape=(batch_size, units)) + else: + c_states = initial_internal_states[0] + h_states = initial_internal_states[1] + + # Create a placeholder for all n-time step outputs. + if time_major: + unrolled_outputs = np.zeros(shape=(sequence_length, batch_size, units)) + else: + unrolled_outputs = np.zeros(shape=(batch_size, sequence_length, units)) + + # Push the batch 4 times through the LSTM cell and capture the outputs plus + # the final h- and c-states. + for t in range(sequence_length): + input_matrix = x[t, :, :] if time_major else x[:, t, :] + input_matrix = np.concatenate((input_matrix, h_states), axis=1) + input_matmul_matrix = np.matmul(input_matrix, weights) + biases + # Forget gate (3rd slot in tf output matrix). Add static forget bias. + sigmoid_1 = sigmoid(input_matmul_matrix[:, units*2:units*3] + + forget_bias) + c_states = np.multiply(c_states, sigmoid_1) + # Add gate (1st and 2nd slots in tf output matrix). + sigmoid_2 = sigmoid(input_matmul_matrix[:, 0:units]) + tanh_3 = np.tanh(input_matmul_matrix[:, units:units*2]) + c_states = np.add(c_states, np.multiply(sigmoid_2, tanh_3)) + # Output gate (last slot in tf output matrix). + sigmoid_4 = sigmoid(input_matmul_matrix[:, units*3:units*4]) + h_states = np.multiply(sigmoid_4, np.tanh(c_states)) + + # Store this output time-slice. + if time_major: + unrolled_outputs[t, :, :] = h_states + else: + unrolled_outputs[:, t, :] = h_states + + return unrolled_outputs, (c_states, h_states) diff --git a/rllib/utils/test_utils.py b/rllib/utils/test_utils.py new file mode 100644 index 000000000..04078702c --- /dev/null +++ b/rllib/utils/test_utils.py @@ -0,0 +1,112 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from ray.rllib.utils.framework import try_import_tf + +tf = try_import_tf() + + +def check(x, y, decimals=5, atol=None, rtol=None, false=False): + """ + Checks two structures (dict, tuple, list, + np.array, float, int, etc..) for (almost) numeric identity. + All numbers in the two structures have to match up to `decimal` digits + after the floating point. Uses assertions. + + Args: + x (any): The first value to be compared (to `y`). + y (any): The second value to be compared (to `x`). + decimals (int): The number of digits after the floating point up to + which all numeric values have to match. + atol (float): Absolute tolerance of the difference between x and y + (overrides `decimals` if given). + rtol (float): Relative tolerance of the difference between x and y + (overrides `decimals` if given). + false (bool): Whether to check that x and y are NOT the same. + """ + # A dict type. + if isinstance(x, dict): + assert isinstance(y, dict), \ + "ERROR: If x is dict, y needs to be a dict as well!" + y_keys = set(x.keys()) + for key, value in x.items(): + assert key in y, \ + "ERROR: y does not have x's key='{}'! y={}".format(key, y) + check(value, y[key], decimals=decimals, atol=atol, rtol=rtol, + false=false) + y_keys.remove(key) + assert not y_keys, \ + "ERROR: y contains keys ({}) that are not in x! y={}".\ + format(list(y_keys), y) + # A tuple type. + elif isinstance(x, (tuple, list)): + assert isinstance(y, (tuple, list)),\ + "ERROR: If x is tuple, y needs to be a tuple as well!" + assert len(y) == len(x),\ + "ERROR: y does not have the same length as x ({} vs {})!".\ + format(len(y), len(x)) + for i, value in enumerate(x): + check(value, y[i], decimals=decimals, atol=atol, rtol=rtol, + false=false) + # Boolean comparison. + elif isinstance(x, (np.bool_, bool)): + if false is True: + assert bool(x) is not bool(y), \ + "ERROR: x ({}) is y ({})!".format(x, y) + else: + assert bool(x) is bool(y), \ + "ERROR: x ({}) is not y ({})!".format(x, y) + # Nones. + elif x is None or y is None: + if false is True: + assert x != y, "ERROR: x ({}) is the same as y ({})!".format(x, y) + else: + assert x == y, \ + "ERROR: x ({}) is not the same as y ({})!".format(x, y) + # String comparison. + elif hasattr(x, "dtype") and x.dtype == np.object: + try: + np.testing.assert_array_equal(x, y) + if false is True: + assert False, \ + "ERROR: x ({}) is the same as y ({})!".format(x, y) + except AssertionError as e: + if false is False: + raise e + # Everything else (assume numeric). + else: + # Numpyize tensors if necessary. + if tf is not None and isinstance(x, tf.Tensor): + x = x.numpy() + if tf is not None and isinstance(y, tf.Tensor): + y = y.numpy() + + # Using decimals. + if atol is None and rtol is None: + try: + np.testing.assert_almost_equal(x, y, decimal=decimals) + if false is True: + assert False, \ + "ERROR: x ({}) is the same as y ({})!".format(x, y) + except AssertionError as e: + if false is False: + raise e + + # Using atol/rtol. + else: + # Provide defaults for either one of atol/rtol. + if atol is None: + atol = 0 + if rtol is None: + rtol = 1e-7 + try: + np.testing.assert_allclose(x, y, atol=atol, rtol=rtol) + if false is True: + assert False, \ + "ERROR: x ({}) is the same as y ({})!".format(x, y) + except AssertionError as e: + if false is False: + raise e diff --git a/rllib/utils/timer.py b/rllib/utils/timer.py index 5dc4e328c..bfad20aab 100644 --- a/rllib/utils/timer.py +++ b/rllib/utils/timer.py @@ -2,9 +2,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import time - import numpy as np +import time class TimerStat(object): diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py new file mode 100644 index 000000000..91a952f90 --- /dev/null +++ b/rllib/utils/torch_ops.py @@ -0,0 +1,23 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from ray.rllib.utils.framework import try_import_torch + +torch, _ = try_import_torch() + + +def sequence_mask(lengths, maxlen, dtype=torch.bool): + """ + Exact same behavior as tf.sequence_mask. + Thanks to Dimitris Papatheodorou + (https://discuss.pytorch.org/t/pytorch-equivalent-for-tf-sequence-mask/39036). + """ + if maxlen is None: + maxlen = lengths.max() + + mask = ~(torch.ones((len(lengths), maxlen)).cumsum(dim=1).t() > lengths).\ + t() + mask.type(dtype) + + return mask