From f9621ce23ca8e9337fb86edfb783f6a57133c604 Mon Sep 17 00:00:00 2001
From: desktable <desktable@users.noreply.github.com>
Date: Wed, 7 Oct 2020 23:23:27 -0700
Subject: [PATCH] [RLlib] Add recsim_wrapper unit test to BUILD (#11225)

---
 ci/travis/install-dependencies.sh    | 2 ++
 rllib/BUILD                          | 7 +++++++
 rllib/env/wrappers/recsim_wrapper.py | 7 ++++---
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/ci/travis/install-dependencies.sh b/ci/travis/install-dependencies.sh
index b0d10be6c..f3b4079fd 100755
--- a/ci/travis/install-dependencies.sh
+++ b/ci/travis/install-dependencies.sh
@@ -286,6 +286,8 @@ install_dependencies() {
   # Additional RLlib dependencies.
   if [ "${RLLIB_TESTING-}" = 1 ]; then
     pip install -r "${WORKSPACE_DIR}"/python/requirements_rllib.txt
+    # install the following packages for testing on travis only
+    pip install 'recsim>=0.2.4'
   fi
 
   # Additional Tune test dependencies.
diff --git a/rllib/BUILD b/rllib/BUILD
index e4a7c5096..7e48f70e2 100644
--- a/rllib/BUILD
+++ b/rllib/BUILD
@@ -1026,6 +1026,13 @@ sh_test(
     data = glob(["examples/serving/*.py"]),
 )
 
+py_test(
+    name = "env/wrappers/tests/test_recsim_wrapper",
+    tags = ["env"],
+    size = "small",
+    srcs = ["env/wrappers/tests/test_recsim_wrapper.py"]
+)
+
 # --------------------------------------------------------------------
 # Models and Distributions
 # rllib/models/
diff --git a/rllib/env/wrappers/recsim_wrapper.py b/rllib/env/wrappers/recsim_wrapper.py
index 6446d19fb..cab1db7ea 100644
--- a/rllib/env/wrappers/recsim_wrapper.py
+++ b/rllib/env/wrappers/recsim_wrapper.py
@@ -4,10 +4,11 @@ RecSim is a configurable recommender systems simulation platform.
 Source: https://github.com/google-research/recsim
 """
 
-import math
-from typing import List, OrderedDict
+from collections import OrderedDict
+from typing import List
 
 import gym
+import numpy as np
 from gym import spaces
 from recsim.environments import interest_evolution
 
@@ -83,7 +84,7 @@ class MultiDiscreteToDiscreteActionWrapper(gym.ActionWrapper):
                 f"is not supported by {self.__class__.__name__}")
         self.action_space_dimensions = env.action_space.nvec
         self.action_space = spaces.Discrete(
-            math.prod(self.action_space_dimensions))
+            np.prod(self.action_space_dimensions))
 
     def action(self, action: int) -> List[int]:
         """Convert a Discrete action to a MultiDiscrete action"""