From 124c8318a876194f25d0ab2382502a3662925b04 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Fri, 18 Dec 2020 00:44:26 +0100 Subject: [PATCH] [RLlib] Fix broken test_distributions.py (test_categorical) (#12915) --- rllib/BUILD | 13 ++++++------- rllib/models/tests/test_distributions.py | 5 +++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/rllib/BUILD b/rllib/BUILD index bde7d5855..bd612d0ff 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1089,13 +1089,12 @@ py_test( srcs = ["models/tests/test_convtranspose2d_stack.py"] ) -# Failing after the following PR: https://github.com/ray-project/ray/pull/12760. -#py_test( -# name = "test_distributions", -# tags = ["models"], -# size = "medium", -# srcs = ["models/tests/test_distributions.py"] -#) +py_test( + name = "test_distributions", + tags = ["models"], + size = "medium", + srcs = ["models/tests/test_distributions.py"] +) # -------------------------------------------------------------------- # Evaluation components diff --git a/rllib/models/tests/test_distributions.py b/rllib/models/tests/test_distributions.py index e0317cf11..987f76a56 100644 --- a/rllib/models/tests/test_distributions.py +++ b/rllib/models/tests/test_distributions.py @@ -87,14 +87,15 @@ class TestDistributions(unittest.TestCase): batch_size = 10000 num_categories = 4 # Create categorical distribution with n categories. - inputs_space = Box(-1.0, 2.0, shape=(batch_size, num_categories)) + inputs_space = Box( + -1.0, 2.0, shape=(batch_size, num_categories), dtype=np.float32) 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, frameworks=("jax", "tf", "tf2", "torch")): + session=True, frameworks=("tf", "tf2", "torch")): # Create the correct distribution object. cls = JAXCategorical if fw == "jax" else Categorical if \ fw != "torch" else TorchCategorical