def test_dist_categorical(self):
dist_sample_shape_correct = [1]
dist_log_probs_correct = [-2.30259]
dist = Categorical([0.1, 0.2, 0.7])
dist_sample_shape = list(dist.sample().size())
dist_log_probs = util.to_numpy(dist.log_prob(0))
util.debug('dist_sample_shape', 'dist_sample_shape_correct', 'dist_log_probs', 'dist_log_probs_correct')
self.assertEqual(dist_sample_shape, dist_sample_shape_correct)
self.assertTrue(np.allclose(dist_log_probs, dist_log_probs_correct, atol=0.1))
def test_dist_categorical_batched(self):
dist_sample_shape_correct = [2]
dist_log_probs_correct = [[-2.30259], [-0.693147]]
dist = Categorical([[0.1, 0.2, 0.7],
[0.2, 0.5, 0.3]])
dist_sample_shape = list(dist.sample().size())
dist_log_probs = util.to_numpy(dist.log_prob([[0, 1]]))
util.debug('dist_sample_shape', 'dist_sample_shape_correct', 'dist_log_probs', 'dist_log_probs_correct')
self.assertEqual(dist_sample_shape, dist_sample_shape_correct)
self.assertTrue(np.allclose(dist_log_probs, dist_log_probs_correct, atol=0.1))