From 8a891b3c304e4043dd27409d993582290da7c439 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Sat, 5 Sep 2020 22:26:42 +0200 Subject: [PATCH] [RLlib] SAC n_step > 1. (#10567) --- rllib/agents/dqn/dqn_tf_policy.py | 4 +- rllib/agents/sac/sac_tf_policy.py | 2 - rllib/agents/sac/sac_torch_policy.py | 2 - rllib/models/tests/test_distributions.py | 91 +++++++++++++++++----- rllib/tuned_examples/sac/pendulum-sac.yaml | 2 +- 5 files changed, 76 insertions(+), 25 deletions(-) diff --git a/rllib/agents/dqn/dqn_tf_policy.py b/rllib/agents/dqn/dqn_tf_policy.py index b2a30088a..ddce5b332 100644 --- a/rllib/agents/dqn/dqn_tf_policy.py +++ b/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], diff --git a/rllib/agents/sac/sac_tf_policy.py b/rllib/agents/sac/sac_tf_policy.py index 987b73ced..26ff292d7 100644 --- a/rllib/agents/sac/sac_tf_policy.py +++ b/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] + diff --git a/rllib/agents/sac/sac_torch_policy.py b/rllib/agents/sac/sac_torch_policy.py index 4e405e0fc..7d9b748fd 100644 --- a/rllib/agents/sac/sac_torch_policy.py +++ b/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] + diff --git a/rllib/models/tests/test_distributions.py b/rllib/models/tests/test_distributions.py index 3c1845919..93e61c30b 100644 --- a/rllib/models/tests/test_distributions.py +++ b/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 diff --git a/rllib/tuned_examples/sac/pendulum-sac.yaml b/rllib/tuned_examples/sac/pendulum-sac.yaml index 01c53f98a..138940b1b 100644 --- a/rllib/tuned_examples/sac/pendulum-sac.yaml +++ b/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