[RLlib] SAC n_step > 1. (#10567)

rllib/agents/dqn/dqn_tf_policy.py

@@ -362,7 +362,7 @@ def _adjust_nstep(n_step, gamma, obs, actions, rewards, new_obs, dones):
 
 
 def postprocess_nstep_and_prio(policy, batch, other_agent=None, episode=None):
-    # N-step Q adjustments
+    # N-step Q adjustments.
     if policy.config["n_step"] > 1:
         _adjust_nstep(policy.config["n_step"], policy.config["gamma"],
                       batch[SampleBatch.CUR_OBS], batch[SampleBatch.ACTIONS],
@@ -372,7 +372,7 @@ def postprocess_nstep_and_prio(policy, batch, other_agent=None, episode=None):
     if PRIO_WEIGHTS not in batch:
         batch[PRIO_WEIGHTS] = np.ones_like(batch[SampleBatch.REWARDS])
 
-    # Prioritize on the worker side
+    # Prioritize on the worker side.
     if batch.count > 0 and policy.config["worker_side_prioritization"]:
         td_errors = policy.compute_td_error(
             batch[SampleBatch.CUR_OBS], batch[SampleBatch.ACTIONS],

rllib/agents/sac/sac_tf_policy.py

@@ -211,8 +211,6 @@ def sac_actor_critic_loss(policy, model, _, train_batch):
         q_tp1_best_masked = (1.0 - tf.cast(train_batch[SampleBatch.DONES],
                                            tf.float32)) * q_tp1_best
 
-    assert policy.config["n_step"] == 1, "TODO(hartikainen) n_step > 1"
-
     # compute RHS of bellman equation
     q_t_selected_target = tf.stop_gradient(
         train_batch[SampleBatch.REWARDS] +

rllib/agents/sac/sac_torch_policy.py

@@ -153,8 +153,6 @@ def actor_critic_loss(policy, model, _, train_batch):
         q_tp1_best_masked = (1.0 - train_batch[SampleBatch.DONES].float()) * \
             q_tp1_best
 
-    assert policy.config["n_step"] == 1, "TODO(hartikainen) n_step > 1"
-
     # compute RHS of bellman equation
     q_t_selected_target = (
         train_batch[SampleBatch.REWARDS] +

rllib/models/tests/test_distributions.py

@@ -28,7 +28,8 @@ class TestDistributions(unittest.TestCase):
                         network_output_shape,
                         fw,
                         sess=None,
-                        bounds=None):
+                        bounds=None,
+                        extra_kwargs=None):
         extreme_values = [
             0.0,
             float(LARGE_INTEGER),
@@ -51,7 +52,7 @@ class TestDistributions(unittest.TestCase):
                 inputs[batch_item][num] = np.log(
                     max(1, np.random.choice((extreme_values))))
 
-        dist = distribution_cls(inputs, {})
+        dist = distribution_cls(inputs, {}, **(extra_kwargs or {}))
         for _ in range(100):
             sample = dist.sample()
             if fw != "tf":
@@ -63,6 +64,12 @@ class TestDistributions(unittest.TestCase):
             if bounds:
                 assert np.min(sample_check) >= bounds[0]
                 assert np.max(sample_check) <= bounds[1]
+                # Make sure bounds make sense and are actually also being
+                # sampled.
+                if isinstance(bounds[0], int):
+                    assert isinstance(bounds[1], int)
+                    assert bounds[0] in sample_check
+                    assert bounds[1] in sample_check
             logp = dist.logp(sample)
             if fw != "tf":
                 logp_check = logp.numpy()
@@ -72,21 +79,59 @@ class TestDistributions(unittest.TestCase):
             assert np.all(np.isfinite(logp_check))
 
     def test_categorical(self):
-        """Tests the Categorical ActionDistribution (tf only)."""
-        num_samples = 100000
-        logits = tf1.placeholder(tf.float32, shape=(None, 10))
-        z = 8 * (np.random.rand(10) - 0.5)
-        data = np.tile(z, (num_samples, 1))
-        c = Categorical(logits, {})  # dummy config dict
-        sample_op = c.sample()
-        sess = tf1.Session()
-        sess.run(tf1.global_variables_initializer())
-        samples = sess.run(sample_op, feed_dict={logits: data})
-        counts = np.zeros(10)
-        for sample in samples:
-            counts[sample] += 1.0
-        probs = np.exp(z) / np.sum(np.exp(z))
-        self.assertTrue(np.sum(np.abs(probs - counts / num_samples)) <= 0.01)
+        batch_size = 10000
+        num_categories = 4
+        # Create categorical distribution with n categories.
+        inputs_space = Box(-1.0, 2.0, shape=(batch_size, num_categories))
+        values_space = Box(
+            0, num_categories - 1, shape=(batch_size, ), dtype=np.int32)
+
+        inputs = inputs_space.sample()
+
+        for fw, sess in framework_iterator(session=True):
+            # Create the correct distribution object.
+            cls = Categorical if fw != "torch" else TorchCategorical
+            categorical = cls(inputs, {})
+
+            # Do a stability test using extreme NN outputs to see whether
+            # sampling and logp'ing result in NaN or +/-inf values.
+            self._stability_test(
+                cls,
+                inputs_space.shape,
+                fw=fw,
+                sess=sess,
+                bounds=(0, num_categories - 1))
+
+            # Batch of size=3 and deterministic (True).
+            expected = np.transpose(np.argmax(inputs, axis=-1))
+            # Sample, expect always max value
+            # (max likelihood for deterministic draw).
+            out = categorical.deterministic_sample()
+            check(out, expected)
+
+            # Batch of size=3 and non-deterministic -> expect roughly the mean.
+            out = categorical.sample()
+            check(
+                tf.reduce_mean(out)
+                if fw != "torch" else torch.mean(out.float()),
+                1.0,
+                decimals=0)
+
+            # Test log-likelihood outputs.
+            probs = softmax(inputs)
+            values = values_space.sample()
+
+            out = categorical.logp(values
+                                   if fw != "torch" else torch.Tensor(values))
+            expected = []
+            for i in range(batch_size):
+                expected.append(np.sum(np.log(np.array(probs[i][values[i]]))))
+            check(out, expected, decimals=4)
+
+            # Test entropy outputs.
+            out = categorical.entropy()
+            expected_entropy = -np.sum(probs * np.log(probs), -1)
+            check(out, expected_entropy)
 
     def test_multi_categorical(self):
         batch_size = 100
@@ -107,11 +152,21 @@ class TestDistributions(unittest.TestCase):
         input_lengths = [num_categories] * num_sub_distributions
         inputs_split = np.split(inputs, num_sub_distributions, axis=1)
 
-        for fw in framework_iterator():
+        for fw, sess in framework_iterator(session=True):
             # Create the correct distribution object.
             cls = MultiCategorical if fw != "torch" else TorchMultiCategorical
             multi_categorical = cls(inputs, None, input_lengths)
 
+            # Do a stability test using extreme NN outputs to see whether
+            # sampling and logp'ing result in NaN or +/-inf values.
+            self._stability_test(
+                cls,
+                inputs_space.shape,
+                fw=fw,
+                sess=sess,
+                bounds=(0, num_categories - 1),
+                extra_kwargs={"input_lens": input_lengths})
+
             # Batch of size=3 and deterministic (True).
             expected = np.transpose(np.argmax(inputs_split, axis=-1))
             # Sample, expect always max value

rllib/tuned_examples/sac/pendulum-sac.yaml

@@ -20,7 +20,7 @@ pendulum-sac:
         tau: 0.005
         target_entropy: auto
         no_done_at_end: true
-        n_step: 1
+        n_step: 3
         rollout_fragment_length: 1
         prioritized_replay: true
         train_batch_size: 256