[RLlib] Add recsim_wrapper unit test to BUILD (#11225)

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.

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/

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"""