[RLlib] Implement DQN PyTorch distributional head. (#9589)

rllib/BUILD

@@ -1171,7 +1171,7 @@ py_test(
     name = "test_env_with_subprocess",
     main = "tests/test_env_with_subprocess.py",
     tags = ["tests_dir", "tests_dir_E"],
-    size = "small",
+    size = "medium",
     srcs = ["tests/test_env_with_subprocess.py"]
 )
 

rllib/agents/dqn/distributional_q_tf_model.py

@@ -49,12 +49,12 @@ class DistributionalQTFModel(TFModelV2):
                 for DDQN. If True, Q-values are calculated as:
                 Q = (A - mean[A]) + V. If False, raw NN output is interpreted
                 as Q-values.
-            num_atoms (int): if >1, enables distributional DQN
-            use_noisy (bool): use noisy nets
-            v_min (float): min value support for distributional DQN
-            v_max (float): max value support for distributional DQN
-            sigma0 (float): initial value of noisy nets
-            add_layer_norm (bool): Add a LayerNorm after each layer..
+            num_atoms (int): If >1, enables distributional DQN.
+            use_noisy (bool): Use noisy nets.
+            v_min (float): Min value support for distributional DQN.
+            v_max (float): Max value support for distributional DQN.
+            sigma0 (float): Initial value of noisy layers.
+            add_layer_norm (bool): Enable layer norm (for param noise).
 
         Note that the core layers for forward() are not defined here, this
         only defines the layers for the Q head. Those layers for forward()

rllib/agents/dqn/dqn_tf_policy.py

@@ -231,8 +231,8 @@ def build_q_losses(policy, model, _, train_batch):
                 train_batch[SampleBatch.NEXT_OBS],
                 explore=False)
         q_tp1_best_using_online_net = tf.argmax(q_tp1_using_online_net, 1)
-        q_tp1_best_one_hot_selection = tf.one_hot(q_tp1_best_using_online_net,
-                                                  policy.action_space.n)
+        q_tp1_best_one_hot_selection = tf.one_hot(
+            q_tp1_best_using_online_net, policy.action_space.n)
         q_tp1_best = tf.reduce_sum(q_tp1 * q_tp1_best_one_hot_selection, 1)
         q_dist_tp1_best = tf.reduce_sum(
             q_dist_tp1 * tf.expand_dims(q_tp1_best_one_hot_selection, -1), 1)
@@ -246,9 +246,9 @@ def build_q_losses(policy, model, _, train_batch):
     policy.q_loss = QLoss(
         q_t_selected, q_logits_t_selected, q_tp1_best, q_dist_tp1_best,
         train_batch[PRIO_WEIGHTS], train_batch[SampleBatch.REWARDS],
-        tf.cast(train_batch[SampleBatch.DONES],
-                tf.float32), config["gamma"], config["n_step"],
-        config["num_atoms"], config["v_min"], config["v_max"])
+        tf.cast(train_batch[SampleBatch.DONES], tf.float32), config["gamma"],
+        config["n_step"], config["num_atoms"],
+        config["v_min"], config["v_max"])
 
     return policy.q_loss.loss
 

rllib/agents/dqn/dqn_torch_model.py

@@ -1,5 +1,5 @@
-import numpy as np
-
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.modules.noisy_layer import NoisyLayer
 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.utils.framework import try_import_torch
 
@@ -18,10 +18,13 @@ class DQNTorchModel(TorchModelV2, nn.Module):
             model_config,
             name,
             *,
-            dueling=False,
             q_hiddens=(256, ),
+            dueling=False,
             dueling_activation="relu",
+            num_atoms=1,
             use_noisy=False,
+            v_min=-10.0,
+            v_max=10.0,
             sigma0=0.5,
             # TODO(sven): Move `add_layer_norm` into ModelCatalog as
             #  generic option, then error if we use ParameterNoise as
@@ -31,19 +34,22 @@ class DQNTorchModel(TorchModelV2, nn.Module):
         """Initialize variables of this model.
 
         Extra model kwargs:
-            dueling (bool): Whether to build the advantage(A)/value(V) heads
-                for DDQN. If True, Q-values are calculated as:
-                Q = (A - mean[A]) + V. If False, raw NN output is interpreted
-                as Q-values.
             q_hiddens (List[int]): List of layer-sizes after(!) the
                 Advantages(A)/Value(V)-split. Hence, each of the A- and V-
                 branches will have this structure of Dense layers. To define
                 the NN before this A/V-split, use - as always -
                 config["model"]["fcnet_hiddens"].
+            dueling (bool): Whether to build the advantage(A)/value(V) heads
+                for DDQN. If True, Q-values are calculated as:
+                Q = (A - mean[A]) + V. If False, raw NN output is interpreted
+                as Q-values.
             dueling_activation (str): The activation to use for all dueling
                 layers (A- and V-branch). One of "relu", "tanh", "linear".
-            use_noisy (bool): use noisy nets
-            sigma0 (float): initial value of noisy nets
+            num_atoms (int): If >1, enables distributional DQN.
+            use_noisy (bool): Use noisy layers.
+            v_min (float): Min value support for distributional DQN.
+            v_max (float): Max value support for distributional DQN.
+            sigma0 (float): Initial value of noisy layers.
             add_layer_norm (bool): Enable layer norm (for param noise).
         """
         nn.Module.__init__(self)
@@ -51,6 +57,10 @@ class DQNTorchModel(TorchModelV2, nn.Module):
                                             num_outputs, model_config, name)
 
         self.dueling = dueling
+        self.num_atoms = num_atoms
+        self.v_min = v_min
+        self.v_max = v_max
+        self.sigma0 = sigma0
         ins = num_outputs
 
         advantage_module = nn.Sequential()
@@ -58,109 +68,87 @@ class DQNTorchModel(TorchModelV2, nn.Module):
 
         # Dueling case: Build the shared (advantages and value) fc-network.
         for i, n in enumerate(q_hiddens):
-            advantage_module.add_module("dueling_A_{}".format(i),
-                                        nn.Linear(ins, n))
-            value_module.add_module("dueling_V_{}".format(i),
-                                    nn.Linear(ins, n))
-            # Add activations if necessary.
-            if dueling_activation == "relu":
-                advantage_module.add_module("dueling_A_act_{}".format(i),
-                                            nn.ReLU())
-                value_module.add_module("dueling_V_act_{}".format(i),
-                                        nn.ReLU())
-            elif dueling_activation == "tanh":
-                advantage_module.add_module("dueling_A_act_{}".format(i),
-                                            nn.Tanh())
-                value_module.add_module("dueling_V_act_{}".format(i),
-                                        nn.Tanh())
-
-            # Add LayerNorm after each Dense.
-            if add_layer_norm:
-                advantage_module.add_module("LayerNorm_A_{}".format(i),
-                                            nn.LayerNorm(n))
-                value_module.add_module("LayerNorm_V_{}".format(i),
-                                        nn.LayerNorm(n))
+            if use_noisy:
+                advantage_module.add_module(
+                    "dueling_A_{}".format(i),
+                    NoisyLayer(
+                        ins, n, sigma0=self.sigma0,
+                        activation=dueling_activation))
+                value_module.add_module(
+                    "dueling_V_{}".format(i),
+                    NoisyLayer(
+                        ins, n, sigma0=self.sigma0,
+                        activation=dueling_activation))
+            else:
+                advantage_module.add_module(
+                    "dueling_A_{}".format(i),
+                    SlimFC(ins, n, activation_fn=dueling_activation))
+                value_module.add_module(
+                    "dueling_V_{}".format(i),
+                    SlimFC(ins, n, activation_fn=dueling_activation))
+                # Add LayerNorm after each Dense.
+                if add_layer_norm:
+                    advantage_module.add_module(
+                        "LayerNorm_A_{}".format(i), nn.LayerNorm(n))
+                    value_module.add_module(
+                        "LayerNorm_V_{}".format(i), nn.LayerNorm(n))
             ins = n
 
         # Actual Advantages layer (nodes=num-actions).
-        if q_hiddens:
-            advantage_module.add_module("A", nn.Linear(ins, action_space.n))
+        if use_noisy:
+            advantage_module.add_module("A", NoisyLayer(
+                ins,
+                self.action_space.n * self.num_atoms,
+                sigma0,
+                activation=None))
+        elif q_hiddens:
+            advantage_module.add_module(
+                "A",
+                SlimFC(
+                    ins, action_space.n * self.num_atoms,
+                    activation_fn=None))
 
         self.advantage_module = advantage_module
 
         # Value layer (nodes=1).
         if self.dueling:
-            value_module.add_module("V", nn.Linear(ins, 1))
+            value_module.add_module("V", SlimFC(ins, 1, activation_fn=None))
             self.value_module = value_module
 
-    def get_advantages_or_q_values(self, model_out):
+    def get_q_value_distributions(self, model_out):
         """Returns distributional values for Q(s, a) given a state embedding.
 
         Override this in your custom model to customize the Q output head.
 
-        Arguments:
-            model_out (Tensor): embedding from the model layers
+        Args:
+            model_out (Tensor): Embedding from the model layers.
 
         Returns:
             (action_scores, logits, dist) if num_atoms == 1, otherwise
             (action_scores, z, support_logits_per_action, logits, dist)
         """
+        action_scores = self.advantage_module(model_out)
 
-        return self.advantage_module(model_out)
+        if self.num_atoms > 1:
+            # Distributional Q-learning uses a discrete support z
+            # to represent the action value distribution
+            z = torch.range(0.0, self.num_atoms - 1, dtype=torch.float32)
+            z = self.v_min + \
+                z * (self.v_max - self.v_min) / float(self.num_atoms - 1)
+
+            support_logits_per_action = torch.reshape(
+                action_scores, shape=(-1, self.action_space.n, self.num_atoms))
+            support_prob_per_action = nn.functional.softmax(
+                support_logits_per_action)
+            action_scores = torch.sum(z * support_prob_per_action, dim=-1)
+            logits = support_logits_per_action
+            probs = support_prob_per_action
+            return action_scores, z, support_logits_per_action, logits, probs
+        else:
+            logits = torch.unsqueeze(torch.ones_like(action_scores), -1)
+            return action_scores, logits, logits
 
     def get_state_value(self, model_out):
         """Returns the state value prediction for the given state embedding."""
 
         return self.value_module(model_out)
-
-    def _noisy_layer(self, action_in, out_size, sigma0, non_linear=True):
-        """
-        a common dense layer: y = w^{T}x + b
-        a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x +
-            (b+\\epsilon_b*\\sigma_b)
-        where \epsilon are random variables sampled from factorized normal
-        distributions and \\sigma are trainable variables which are expected to
-        vanish along the training procedure
-        """
-        in_size = int(action_in.shape[1])
-
-        epsilon_in = torch.normal(
-            mean=torch.zeros([in_size]), std=torch.ones([in_size]))
-        epsilon_out = torch.normal(
-            mean=torch.zeros([out_size]), std=torch.ones([out_size]))
-        epsilon_in = self._f_epsilon(epsilon_in)
-        epsilon_out = self._f_epsilon(epsilon_out)
-        epsilon_w = torch.matmul(
-            torch.unsqueeze(epsilon_in, -1),
-            other=torch.unsqueeze(epsilon_out, 0))
-        epsilon_b = epsilon_out
-
-        sigma_w = torch.Tensor(
-            data=np.random.uniform(
-                low=-1.0 / np.sqrt(float(in_size)),
-                high=1.0 / np.sqrt(float(in_size)),
-                size=[in_size, out_size]),
-            dtype=torch.float32,
-            requires_grad=True)
-        # TF noise generation can be unreliable on GPU
-        # If generating the noise on the CPU,
-        # lowering sigma0 to 0.1 may be helpful
-        sigma_b = torch.Tensor(
-            data=np.full(
-                shape=[out_size], fill_value=sigma0 / np.sqrt(float(in_size))),
-            requires_grad=True)
-        w = torch.Tensor(
-            data=np.full(
-                shape=[in_size, out_size],
-                fill_value=6 / np.sqrt(float(in_size) + float(out_size))),
-            requires_grad=True)
-        b = torch.Tensor(data=np.zeros([out_size]), requires_grad=True)
-        action_activation = torch.matmul(action_in, w + sigma_w * epsilon_w) \
-            + b + sigma_b * epsilon_b
-
-        if not non_linear:
-            return action_activation
-        return nn.functional.relu(action_activation)
-
-    def _f_epsilon(self, x):
-        return torch.sign(x) * torch.pow(torch.abs(x), 0.5)

rllib/agents/dqn/dqn_torch_policy.py

@@ -14,7 +14,8 @@ from ray.rllib.policy.torch_policy_template import build_torch_policy
 from ray.rllib.utils.error import UnsupportedSpaceException
 from ray.rllib.utils.exploration.parameter_noise import ParameterNoise
 from ray.rllib.utils.framework import try_import_torch
-from ray.rllib.utils.torch_ops import huber_loss, reduce_mean_ignore_inf
+from ray.rllib.utils.torch_ops import huber_loss, reduce_mean_ignore_inf, \
+    softmax_cross_entropy_with_logits
 
 torch, nn = try_import_torch()
 F = None
@@ -25,7 +26,9 @@ if nn:
 class QLoss:
     def __init__(self,
                  q_t_selected,
+                 q_logits_t_selected,
                  q_tp1_best,
+                 q_probs_tp1_best,
                  importance_weights,
                  rewards,
                  done_mask,
@@ -36,24 +39,61 @@ class QLoss:
                  v_max=10.0):
 
         if num_atoms > 1:
-            raise ValueError("Torch version of DQN does not support "
-                             "distributional Q yet!")
+            # Distributional Q-learning which corresponds to an entropy loss
+            z = torch.range(0.0, num_atoms - 1, dtype=torch.float32)
+            z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
 
-        q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
+            # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)
+            r_tau = torch.unsqueeze(
+                rewards, -1) + gamma**n_step * torch.unsqueeze(
+                    1.0 - done_mask, -1) * torch.unsqueeze(z, 0)
+            r_tau = torch.clamp(r_tau, v_min, v_max)
+            b = (r_tau - v_min) / ((v_max - v_min) / float(num_atoms - 1))
+            lb = torch.floor(b)
+            ub = torch.ceil(b)
 
-        # compute RHS of bellman equation
-        q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked
+            # Indispensable judgement which is missed in most implementations
+            # when b happens to be an integer, lb == ub, so pr_j(s', a*) will
+            # be discarded because (ub-b) == (b-lb) == 0.
+            floor_equal_ceil = (ub - lb < 0.5).float()
 
-        # compute the error (potentially clipped)
-        self.td_error = q_t_selected - q_t_selected_target.detach()
-        self.loss = torch.mean(
-            importance_weights.float() * huber_loss(self.td_error))
-        self.stats = {
-            "mean_q": torch.mean(q_t_selected),
-            "min_q": torch.min(q_t_selected),
-            "max_q": torch.max(q_t_selected),
-            "mean_td_error": torch.mean(self.td_error),
-        }
+            # (batch_size, num_atoms, num_atoms)
+            l_project = F.one_hot(lb.long(), num_atoms)
+            # (batch_size, num_atoms, num_atoms)
+            u_project = F.one_hot(ub.long(), num_atoms)
+            ml_delta = q_probs_tp1_best * (ub - b + floor_equal_ceil)
+            mu_delta = q_probs_tp1_best * (b - lb)
+            ml_delta = torch.sum(
+                l_project * torch.unsqueeze(ml_delta, -1), dim=1)
+            mu_delta = torch.sum(
+                u_project * torch.unsqueeze(mu_delta, -1), dim=1)
+            m = ml_delta + mu_delta
+
+            # Rainbow paper claims that using this cross entropy loss for
+            # priority is robust and insensitive to `prioritized_replay_alpha`
+            self.td_error = softmax_cross_entropy_with_logits(
+                logits=q_logits_t_selected, labels=m)
+            self.loss = torch.mean(self.td_error * importance_weights)
+            self.stats = {
+                # TODO: better Q stats for dist dqn
+                "mean_td_error": torch.mean(self.td_error),
+            }
+        else:
+            q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
+
+            # compute RHS of bellman equation
+            q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked
+
+            # compute the error (potentially clipped)
+            self.td_error = q_t_selected - q_t_selected_target.detach()
+            self.loss = torch.mean(
+                importance_weights.float() * huber_loss(self.td_error))
+            self.stats = {
+                "mean_q": torch.mean(q_t_selected),
+                "min_q": torch.min(q_t_selected),
+                "max_q": torch.max(q_t_selected),
+                "mean_td_error": torch.mean(self.td_error),
+            }
 
 
 class ComputeTDErrorMixin:
@@ -102,9 +142,12 @@ def build_q_model_and_distribution(policy, obs_space, action_space, config):
         framework="torch",
         model_interface=DQNTorchModel,
         name=Q_SCOPE,
-        dueling=config["dueling"],
         q_hiddens=config["hiddens"],
+        dueling=config["dueling"],
+        num_atoms=config["num_atoms"],
         use_noisy=config["noisy"],
+        v_min=config["v_min"],
+        v_max=config["v_max"],
         sigma0=config["sigma0"],
         # TODO(sven): Move option to add LayerNorm after each Dense
         #  generically into ModelCatalog.
@@ -120,9 +163,12 @@ def build_q_model_and_distribution(policy, obs_space, action_space, config):
         framework="torch",
         model_interface=DQNTorchModel,
         name=Q_TARGET_SCOPE,
-        dueling=config["dueling"],
         q_hiddens=config["hiddens"],
+        dueling=config["dueling"],
+        num_atoms=config["num_atoms"],
         use_noisy=config["noisy"],
+        v_min=config["v_min"],
+        v_max=config["v_max"],
         sigma0=config["sigma0"],
         # TODO(sven): Move option to add LayerNorm after each Dense
         #  generically into ModelCatalog.
@@ -149,54 +195,63 @@ def get_distribution_inputs_and_class(policy,
 
 def build_q_losses(policy, model, _, train_batch):
     config = policy.config
-    # q network evaluation
-    q_t = compute_q_values(
+    # Q-network evaluation.
+    q_t, q_logits_t, q_probs_t = compute_q_values(
         policy,
         policy.q_model,
         train_batch[SampleBatch.CUR_OBS],
         explore=False,
         is_training=True)
 
-    # target q network evalution
-    q_tp1 = compute_q_values(
+    # Target Q-network evaluation.
+    q_tp1, q_logits_tp1, q_probs_tp1 = compute_q_values(
         policy,
         policy.target_q_model,
         train_batch[SampleBatch.NEXT_OBS],
         explore=False,
         is_training=True)
 
-    # q scores for actions which we know were selected in the given state.
-    one_hot_selection = F.one_hot(train_batch[SampleBatch.ACTIONS],
-                                  policy.action_space.n)
+    # Q scores for actions which we know were selected in the given state.
+    one_hot_selection = F.one_hot(
+        train_batch[SampleBatch.ACTIONS], policy.action_space.n)
     q_t_selected = torch.sum(
         torch.where(q_t > -float("inf"), q_t, torch.tensor(0.0)) *
         one_hot_selection, 1)
+    q_logits_t_selected = torch.sum(
+        q_logits_t * torch.unsqueeze(one_hot_selection, -1), 1)
 
     # compute estimate of best possible value starting from state at t + 1
     if config["double_q"]:
-        q_tp1_using_online_net = compute_q_values(
-            policy,
-            policy.q_model,
-            train_batch[SampleBatch.NEXT_OBS],
-            explore=False,
-            is_training=True)
+        q_tp1_using_online_net, q_logits_tp1_using_online_net, \
+            q_dist_tp1_using_online_net = compute_q_values(
+                policy,
+                policy.q_model,
+                train_batch[SampleBatch.NEXT_OBS],
+                explore=False,
+                is_training=True)
         q_tp1_best_using_online_net = torch.argmax(q_tp1_using_online_net, 1)
-        q_tp1_best_one_hot_selection = F.one_hot(q_tp1_best_using_online_net,
-                                                 policy.action_space.n)
+        q_tp1_best_one_hot_selection = F.one_hot(
+            q_tp1_best_using_online_net, policy.action_space.n)
+        q_tp1_best = torch.sum(
+            torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) *
+            q_tp1_best_one_hot_selection, 1)
+        q_probs_tp1_best = torch.sum(
+            q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1)
     else:
         q_tp1_best_one_hot_selection = F.one_hot(
             torch.argmax(q_tp1, 1), policy.action_space.n)
+        q_tp1_best = torch.sum(
+            torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) *
+            q_tp1_best_one_hot_selection, 1)
+        q_probs_tp1_best = torch.sum(
+            q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1)
 
-    q_tp1_best = torch.sum(
-        torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) *
-        q_tp1_best_one_hot_selection, 1)
-
-    policy.q_loss = QLoss(q_t_selected, q_tp1_best, train_batch[PRIO_WEIGHTS],
-                          train_batch[SampleBatch.REWARDS],
-                          train_batch[SampleBatch.DONES].float(),
-                          config["gamma"], config["n_step"],
-                          config["num_atoms"], config["v_min"],
-                          config["v_max"])
+    policy.q_loss = QLoss(
+        q_t_selected, q_logits_t_selected, q_tp1_best, q_probs_tp1_best,
+        train_batch[PRIO_WEIGHTS], train_batch[SampleBatch.REWARDS],
+        train_batch[SampleBatch.DONES].float(), config["gamma"],
+        config["n_step"], config["num_atoms"],
+        config["v_min"], config["v_max"])
 
     return policy.q_loss.loss
 
@@ -225,26 +280,44 @@ def after_init(policy, obs_space, action_space, config):
 
 
 def compute_q_values(policy, model, obs, explore, is_training=False):
-    if policy.config["num_atoms"] > 1:
-        raise ValueError("torch DQN does not support distributional DQN yet!")
+    config = policy.config
 
     model_out, state = model({
         SampleBatch.CUR_OBS: obs,
         "is_training": is_training,
     }, [], None)
 
-    advantages_or_q_values = model.get_advantages_or_q_values(model_out)
-
-    if policy.config["dueling"]:
-        state_value = model.get_state_value(model_out)
-        advantages_mean = reduce_mean_ignore_inf(advantages_or_q_values, 1)
-        advantages_centered = advantages_or_q_values - torch.unsqueeze(
-            advantages_mean, 1)
-        q_values = state_value + advantages_centered
+    if config["num_atoms"] > 1:
+        (action_scores, z, support_logits_per_action, logits,
+         probs_or_logits) = model.get_q_value_distributions(model_out)
     else:
-        q_values = advantages_or_q_values
+        (action_scores, logits,
+         probs_or_logits) = model.get_q_value_distributions(model_out)
 
-    return q_values
+    if config["dueling"]:
+        state_score = model.get_state_value(model_out)
+        if policy.config["num_atoms"] > 1:
+            support_logits_per_action_mean = torch.mean(
+                support_logits_per_action, dim=1)
+            support_logits_per_action_centered = (
+                support_logits_per_action - torch.unsqueeze(
+                    support_logits_per_action_mean, dim=1))
+            support_logits_per_action = torch.unsqueeze(
+                state_score, dim=1) + support_logits_per_action_centered
+            support_prob_per_action = nn.functional.softmax(
+                support_logits_per_action)
+            value = torch.sum(z * support_prob_per_action, dim=-1)
+            logits = support_logits_per_action
+            probs_or_logits = support_prob_per_action
+        else:
+            advantages_mean = reduce_mean_ignore_inf(action_scores, 1)
+            advantages_centered = action_scores - torch.unsqueeze(
+                advantages_mean, 1)
+            value = state_score + advantages_centered
+    else:
+        value = action_scores
+
+    return value, logits, probs_or_logits
 
 
 def grad_process_and_td_error_fn(policy, optimizer, loss):

rllib/agents/dqn/tests/test_dqn.py

@@ -22,8 +22,9 @@ class TestDQN(unittest.TestCase):
         config["num_workers"] = 2
         num_iterations = 1
 
-        for fw in framework_iterator(config):
+        for _ in framework_iterator(config):
             # Double-dueling DQN.
+            print("Double-dueling")
             plain_config = config.copy()
             trainer = dqn.DQNTrainer(config=plain_config, env="CartPole-v0")
             for i in range(num_iterations):
@@ -34,9 +35,7 @@ class TestDQN(unittest.TestCase):
             trainer.stop()
 
             # Rainbow.
-            # TODO(sven): Add torch once DQN-torch supports distributional-Q.
-            if fw == "torch":
-                continue
+            print("Rainbow")
             rainbow_config = config.copy()
             rainbow_config["num_atoms"] = 10
             rainbow_config["noisy"] = True

rllib/examples/env/env_with_subprocess.py

@@ -3,14 +3,13 @@ import gym
 from gym.spaces import Discrete
 import os
 import subprocess
-import time
 
 
 class EnvWithSubprocess(gym.Env):
-    """Our env that spawns a subprocess."""
+    """An env that spawns a subprocess."""
 
     # Dummy command to run as a subprocess with a unique name
-    UNIQUE_CMD = "sleep {}".format(str(time.time()))
+    UNIQUE_CMD = "sleep 20"
 
     def __init__(self, config):
         self.UNIQUE_FILE_0 = config["tmp_file1"]
@@ -20,11 +19,11 @@ class EnvWithSubprocess(gym.Env):
 
         self.action_space = Discrete(2)
         self.observation_space = Discrete(2)
-        # Subprocess that should be cleaned up
+        # Subprocess that should be cleaned up.
         self.subproc = subprocess.Popen(
             self.UNIQUE_CMD.split(" "), shell=False)
         self.config = config
-        # Exit handler should be called
+        # Exit handler should be called.
         atexit.register(lambda: self.subproc.kill())
         if config.worker_index == 0:
             atexit.register(lambda: os.unlink(self.UNIQUE_FILE_0))

rllib/models/torch/misc.py

@@ -1,7 +1,7 @@
 """ Code adapted from https://github.com/ikostrikov/pytorch-a3c"""
 import numpy as np
 
-from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.framework import get_activation_fn, try_import_torch
 
 torch, nn = try_import_torch()
 
@@ -19,17 +19,21 @@ def same_padding(in_size, filter_size, stride_size):
     """Note: Padding is added to match TF conv2d `same` padding. See
     www.tensorflow.org/versions/r0.12/api_docs/python/nn/convolution
 
-    Params:
+    Args:
         in_size (tuple): Rows (Height), Column (Width) for input
-        stride_size (tuple): Rows (Height), Column (Width) for stride
-        filter_size (tuple): Rows (Height), Column (Width) for filter
+        stride_size (Union[int,Tuple[int, int]]): Rows (Height), column (Width)
+            for stride. If int, height == width.
+        filter_size (tuple): Rows (Height), column (Width) for filter
 
-    Output:
+    Returns:
         padding (tuple): For input into torch.nn.ZeroPad2d.
         output (tuple): Output shape after padding and convolution.
     """
     in_height, in_width = in_size
-    filter_height, filter_width = filter_size
+    if isinstance(filter_size, int):
+        filter_height, filter_width = filter_size, filter_size
+    else:
+        filter_height, filter_width = filter_size
     stride_height, stride_width = stride_size
 
     out_height = np.ceil(float(in_height) / float(stride_height))
@@ -102,7 +106,9 @@ class SlimFC(nn.Module):
         if use_bias is True:
             nn.init.constant_(linear.bias, bias_init)
         layers.append(linear)
-        if activation_fn:
+        if isinstance(activation_fn, str):
+            activation_fn = get_activation_fn(activation_fn, "torch")
+        if activation_fn is not None:
             layers.append(activation_fn())
         self._model = nn.Sequential(*layers)
 

rllib/models/torch/modules/noisy_layer.py

@@ -0,0 +1,90 @@
+import numpy as np
+
+from ray.rllib.utils.framework import get_activation_fn, try_import_torch
+from ray.rllib.utils.framework import get_variable
+
+torch, nn = try_import_torch()
+
+
+class NoisyLayer(nn.Module):
+    """A Layer that adds learnable Noise
+    a common dense layer: y = w^{T}x + b
+    a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x +
+        (b+\\epsilon_b*\\sigma_b)
+    where \epsilon are random variables sampled from factorized normal
+    distributions and \\sigma are trainable variables which are expected to
+    vanish along the training procedure
+    """
+
+    def __init__(self, in_size, out_size, sigma0, activation="relu"):
+        """Initializes a NoisyLayer object.
+
+        Args:
+            in_size:
+            out_size:
+            sigma0:
+            non_linear:
+        """
+        super().__init__()
+
+        self.in_size = in_size
+        self.out_size = out_size
+        self.sigma0 = sigma0
+        self.activation = get_activation_fn(activation, framework="torch")
+        if self.activation is not None:
+            self.activation = self.activation()
+
+        self.sigma_w = get_variable(
+            np.random.uniform(
+                low=-1.0 / np.sqrt(float(self.in_size)),
+                high=1.0 / np.sqrt(float(self.in_size)),
+                size=[self.in_size, out_size]),
+            framework="torch",
+            dtype=torch.float32,
+            torch_tensor=True,
+            trainable=True)
+        self.sigma_b = get_variable(
+            np.full(
+                shape=[out_size],
+                fill_value=sigma0 / np.sqrt(float(self.in_size))),
+            framework="torch",
+            dtype=torch.float32,
+            torch_tensor=True,
+            trainable=True)
+        self.w = get_variable(
+            np.full(
+                shape=[self.in_size, self.out_size],
+                fill_value=6 / np.sqrt(float(in_size) + float(out_size))),
+            framework="torch",
+            dtype=torch.float32,
+            torch_tensor=True,
+            trainable=True)
+        self.b = get_variable(
+            np.zeros([out_size]),
+            framework="torch",
+            dtype=torch.float32,
+            torch_tensor=True,
+            trainable=True)
+
+    def forward(self, inputs):
+        epsilon_in = self._f_epsilon(torch.normal(
+            mean=torch.zeros([self.in_size]),
+            std=torch.ones([self.in_size])))
+        epsilon_out = self._f_epsilon(torch.normal(
+            mean=torch.zeros([self.out_size]),
+            std=torch.ones([self.out_size])))
+        epsilon_w = torch.matmul(
+            torch.unsqueeze(epsilon_in, -1),
+            other=torch.unsqueeze(epsilon_out, 0))
+        epsilon_b = epsilon_out
+
+        action_activation = torch.matmul(
+            inputs, self.w + self.sigma_w * epsilon_w
+        ) + self.b + self.sigma_b * epsilon_b
+
+        if self.activation is not None:
+            action_activation = self.activation(action_activation)
+        return action_activation
+
+    def _f_epsilon(self, x):
+        return torch.sign(x) * torch.pow(torch.abs(x), 0.5)

rllib/tests/test_rollout_worker.py

@@ -267,7 +267,7 @@ class TestRolloutWorker(unittest.TestCase):
             pg.stop()
 
     def test_reward_clipping(self):
-        # clipping on
+        # Clipping: on.
         ev = RolloutWorker(
             env_creator=lambda _: MockEnv2(episode_length=10),
             policy=MockPolicy,
@@ -278,7 +278,7 @@ class TestRolloutWorker(unittest.TestCase):
         self.assertEqual(result["episode_reward_mean"], 1000)
         ev.stop()
 
-        # clipping off
+        # Clipping: off.
         ev2 = RolloutWorker(
             env_creator=lambda _: MockEnv2(episode_length=10),
             policy=MockPolicy,

rllib/utils/framework.py

@@ -206,6 +206,10 @@ def get_variable(value,
     elif framework == "torch" and torch_tensor is True:
         torch, _ = try_import_torch()
         var_ = torch.from_numpy(value)
+        if dtype == torch.float32:
+            var_ = var_.float()
+        elif dtype == torch.int32:
+            var_ = var_.int()
         if device:
             var_ = var_.to(device)
         var_.requires_grad = trainable

rllib/utils/torch_ops.py

@@ -4,7 +4,11 @@ import tree
 from ray.rllib.models.repeated_values import RepeatedValues
 from ray.rllib.utils.framework import try_import_torch
 
-torch, _ = try_import_torch()
+torch, nn = try_import_torch()
+
+
+def atanh(x):
+    return 0.5 * torch.log((1 + x) / (1 - x))
 
 
 def explained_variance(y, pred):
@@ -50,13 +54,6 @@ def l2_loss(x):
     return torch.sum(torch.pow(x, 2.0)) / 2.0
 
 
-def reduce_mean_ignore_inf(x, axis):
-    """Same as torch.mean() but ignores -inf values."""
-    mask = torch.ne(x, float("-inf"))
-    x_zeroed = torch.where(mask, x, torch.zeros_like(x))
-    return torch.sum(x_zeroed, axis) / torch.sum(mask.float(), axis)
-
-
 def minimize_and_clip(optimizer, clip_val=10):
     """Clips gradients found in `optimizer.param_groups` to given value.
 
@@ -69,6 +66,13 @@ def minimize_and_clip(optimizer, clip_val=10):
                 torch.nn.utils.clip_grad_norm_(p.grad, clip_val)
 
 
+def reduce_mean_ignore_inf(x, axis):
+    """Same as torch.mean() but ignores -inf values."""
+    mask = torch.ne(x, float("-inf"))
+    x_zeroed = torch.where(mask, x, torch.zeros_like(x))
+    return torch.sum(x_zeroed, axis) / torch.sum(mask.float(), axis)
+
+
 def sequence_mask(lengths, maxlen=None, dtype=None):
     """Offers same behavior as tf.sequence_mask for torch.
 
@@ -86,6 +90,18 @@ def sequence_mask(lengths, maxlen=None, dtype=None):
     return mask
 
 
+def softmax_cross_entropy_with_logits(logits, labels):
+    """Same behavior as tf.nn.softmax_cross_entropy_with_logits.
+
+    Args:
+        x (TensorType):
+
+    Returns:
+
+    """
+    return torch.sum(-labels * nn.functional.log_softmax(logits, -1), -1)
+
+
 def convert_to_non_torch_type(stats):
     """Converts values in `stats` to non-Tensor numpy or python types.
 
@@ -138,7 +154,3 @@ def convert_to_torch_tensor(x, device=None):
         return tensor if device is None else tensor.to(device)
 
     return tree.map_structure(mapping, x)
-
-
-def atanh(x):
-    return 0.5 * torch.log((1 + x) / (1 - x))