diff --git a/rllib/env/wrappers/__init__.py b/rllib/env/wrappers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/env/wrappers/recsim_wrapper.py b/rllib/env/wrappers/recsim_wrapper.py new file mode 100644 index 000000000..6446d19fb --- /dev/null +++ b/rllib/env/wrappers/recsim_wrapper.py @@ -0,0 +1,116 @@ +"""Wrap Google's RecSim environment for RLlib + +RecSim is a configurable recommender systems simulation platform. +Source: https://github.com/google-research/recsim +""" + +import math +from typing import List, OrderedDict + +import gym +from gym import spaces +from recsim.environments import interest_evolution + +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.tune.registry import register_env + + +class RecSimObservationSpaceWrapper(gym.ObservationWrapper): + """Fix RecSim environment's observation space + + In RecSim's observation spaces, the "doc" field is a dictionary keyed by + document IDs. Those IDs are changing every step, thus generating a + different observation space in each time. This causes issues for RLlib + because it expects the observation space to remain the same across steps. + + This environment wrapper fixes that by reindexing the documents by their + positions in the list. + """ + + def __init__(self, env: gym.Env): + super().__init__(env) + obs_space = self.env.observation_space + doc_space = spaces.Dict( + OrderedDict( + [(str(k), doc) + for k, (_, + doc) in enumerate(obs_space["doc"].spaces.items())])) + self.observation_space = spaces.Dict( + OrderedDict([ + ("user", obs_space["user"]), + ("doc", doc_space), + ("response", obs_space["response"]), + ])) + + def observation(self, obs): + new_obs = OrderedDict() + new_obs["user"] = obs["user"] + new_obs["doc"] = { + str(k): v + for k, (_, v) in enumerate(obs["doc"].items()) + } + new_obs["response"] = obs["response"] + return new_obs + + +class RecSimResetWrapper(gym.Wrapper): + """Fix RecSim environment's reset() function + + RecSim's reset() function returns an observation without the "response" + field, breaking RLlib's check. This wrapper fixes that by assigning a + random "response". + """ + + def reset(self): + obs = super().reset() + obs["response"] = self.env.observation_space["response"].sample() + return obs + + +class MultiDiscreteToDiscreteActionWrapper(gym.ActionWrapper): + """Convert the action space from MultiDiscrete to Discrete + + At this moment, RLlib's DQN algorithms only work on Discrete action space. + This wrapper allows us to apply DQN algorithms to the RecSim environment. + """ + + def __init__(self, env: gym.Env): + super().__init__(env) + + if not isinstance(env.action_space, spaces.MultiDiscrete): + raise UnsupportedSpaceException( + f"Action space {env.action_space} " + 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)) + + def action(self, action: int) -> List[int]: + """Convert a Discrete action to a MultiDiscrete action""" + multi_action = [None] * len(self.action_space_dimensions) + for idx, n in enumerate(self.action_space_dimensions): + action, dim_action = divmod(action, n) + multi_action[idx] = dim_action + return multi_action + + +def make_recsim_env(config): + DEFAULT_ENV_CONFIG = { + "num_candidates": 10, + "slate_size": 2, + "resample_documents": True, + "seed": 0, + "convert_to_discrete_action_space": False, + } + env_config = DEFAULT_ENV_CONFIG.copy() + env_config.update(config) + env = interest_evolution.create_environment(env_config) + env = RecSimResetWrapper(env) + env = RecSimObservationSpaceWrapper(env) + if config and config["convert_to_discrete_action_space"]: + env = MultiDiscreteToDiscreteActionWrapper(env) + return env + + +env_name = "RecSim-v1" +register_env(name=env_name, env_creator=make_recsim_env) diff --git a/rllib/env/wrappers/tests/test_recsim_wrapper.py b/rllib/env/wrappers/tests/test_recsim_wrapper.py new file mode 100644 index 000000000..3211dca71 --- /dev/null +++ b/rllib/env/wrappers/tests/test_recsim_wrapper.py @@ -0,0 +1,38 @@ +import unittest + +import gym + +from ray.rllib.env.wrappers.recsim_wrapper import ( + MultiDiscreteToDiscreteActionWrapper, make_recsim_env) +from ray.rllib.utils.error import UnsupportedSpaceException + + +class TestRecSimWrapper(unittest.TestCase): + def test_observation_space(self): + env = make_recsim_env(config={}) + obs = env.reset() + self.assertTrue( + env.observation_space.contains(obs), + f"{env.observation_space} doesn't contain {obs}") + new_obs, _, _, _ = env.step(env.action_space.sample()) + self.assertTrue(env.observation_space.contains(new_obs)) + + def test_action_space_conversion(self): + env = make_recsim_env( + config={"convert_to_discrete_action_space": True}) + self.assertIsInstance(env.action_space, gym.spaces.Discrete) + env.reset() + action = env.action_space.sample() + env.step(action) + + def test_double_action_space_conversion_raises_exception(self): + env = make_recsim_env( + config={"convert_to_discrete_action_space": True}) + with self.assertRaises(UnsupportedSpaceException): + env = MultiDiscreteToDiscreteActionWrapper(env) + + +if __name__ == "__main__": + import sys + import pytest + sys.exit(pytest.main(["-v", __file__]))