Get utils ready for better Agent torch support. (#6561)

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):

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

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):

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 ===

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

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):

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__)
 
 

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):

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()):

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:

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)
 
 

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__)
 

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):

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):

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

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):

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({

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",
 ]

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

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))

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

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)

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

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):

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