From 6ddf84b019be0b6dee52c9699f7c6ebe416dc3d3 Mon Sep 17 00:00:00 2001 From: Saurabh Gupta Date: Thu, 26 Mar 2020 13:41:16 -0700 Subject: [PATCH] Contextual Bandit algorithms (WIP) (#7642) --- .travis.yml | 2 +- doc/source/rllib-algorithms.rst | 63 ++++- rllib/BUILD | 18 ++ rllib/contrib/bandits/__init__.py | 0 rllib/contrib/bandits/agents/__init__.py | 4 + rllib/contrib/bandits/agents/lin_ts.py | 46 ++++ rllib/contrib/bandits/agents/lin_ucb.py | 46 ++++ rllib/contrib/bandits/agents/policy.py | 121 +++++++++ rllib/contrib/bandits/envs/__init__.py | 5 + rllib/contrib/bandits/envs/discrete.py | 171 ++++++++++++ rllib/contrib/bandits/envs/parametric.py | 157 +++++++++++ .../bandits/examples/LinTS_train_wheel_env.py | 52 ++++ .../bandits/examples/simple_context_bandit.py | 47 ++++ .../examples/tune_LinTS_train_wheel_env.py | 80 ++++++ .../tune_LinUCB_train_recommendation.py | 54 ++++ rllib/contrib/bandits/exploration.py | 52 ++++ rllib/contrib/bandits/models/__init__.py | 0 .../bandits/models/linear_regression.py | 246 ++++++++++++++++++ rllib/contrib/registry.py | 12 + 19 files changed, 1174 insertions(+), 2 deletions(-) create mode 100644 rllib/contrib/bandits/__init__.py create mode 100644 rllib/contrib/bandits/agents/__init__.py create mode 100644 rllib/contrib/bandits/agents/lin_ts.py create mode 100644 rllib/contrib/bandits/agents/lin_ucb.py create mode 100644 rllib/contrib/bandits/agents/policy.py create mode 100644 rllib/contrib/bandits/envs/__init__.py create mode 100644 rllib/contrib/bandits/envs/discrete.py create mode 100644 rllib/contrib/bandits/envs/parametric.py create mode 100644 rllib/contrib/bandits/examples/LinTS_train_wheel_env.py create mode 100644 rllib/contrib/bandits/examples/simple_context_bandit.py create mode 100644 rllib/contrib/bandits/examples/tune_LinTS_train_wheel_env.py create mode 100644 rllib/contrib/bandits/examples/tune_LinUCB_train_recommendation.py create mode 100644 rllib/contrib/bandits/exploration.py create mode 100644 rllib/contrib/bandits/models/__init__.py create mode 100644 rllib/contrib/bandits/models/linear_regression.py diff --git a/.travis.yml b/.travis.yml index 4b6043109..421bd1f36 100644 --- a/.travis.yml +++ b/.travis.yml @@ -254,7 +254,7 @@ matrix: - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_A,examples_B --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_C --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_P --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... - - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_R,examples_S,examples_T --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... + - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_U,examples_R,examples_S,examples_T --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... # RLlib: tests_dir: Everything in rllib/tests/ directory (A-I). - os: linux diff --git a/doc/source/rllib-algorithms.rst b/doc/source/rllib-algorithms.rst index 89c42bff7..a7df0b0fd 100644 --- a/doc/source/rllib-algorithms.rst +++ b/doc/source/rllib-algorithms.rst @@ -465,8 +465,69 @@ Tuned examples: `CartPole-v0 `__ `[implementation] +`__ +LinUCB assumes a linear dependency between the expected reward of an action and +its context. It estimates the Q value of each action using ridge regression. +It constructs a confidence region around the weights of the linear +regression model and uses this confidence ellipsoid to estimate the +uncertainty of action values. + +**LinUCB-specific configs** (see also `common configs `__): + +.. literalinclude:: ../../rllib/contrib/bandits/agents/lin_ucb.py + :language: python + :start-after: __sphinx_doc_begin__ + :end-before: __sphinx_doc_end__ + + +LinTS (Linear Thompson Sampling) +-------------------------------- +|pytorch| +`[paper] `__ `[implementation] +`__ +Like LinUCB, LinTS also assumes a linear dependency between the expected +reward of an action and its context and uses online ridge regression to +estimate the Q values of actions given the context. It assumes a Gaussian +prior on the weights and a Gaussian likelihood function. For deciding which +action to take, the agent samples weights for each arm, using +the posterior distributions, and plays the arm that produces the highest reward. + +**LinTS-specific configs** (see also `common configs `__): + +.. literalinclude:: ../../rllib/contrib/bandits/agents/lin_ts.py + :language: python + :start-after: __sphinx_doc_begin__ + :end-before: __sphinx_doc_end__ + + .. |tensorflow| image:: tensorflow.png :width: 24 .. |pytorch| image:: pytorch.png - :width: 24 + :width: 24 \ No newline at end of file diff --git a/rllib/BUILD b/rllib/BUILD index 76bf3b3d8..eea7f3c9b 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1376,6 +1376,24 @@ py_test( args = ["--stop=2000", "--run=contrib/MADDPG"] ) +py_test( + name = "contrib/bandits/examples/lin_ts", + main = "contrib/bandits/examples/simple_context_bandit.py", + tags = ["examples", "examples_T"], + size = "small", + srcs = ["contrib/bandits/examples/simple_context_bandit.py"], + args = ["--stop-at-reward=10", "--run=contrib/LinTS"], +) + +py_test( + name = "contrib/bandits/examples/lin_ucb", + main = "contrib/bandits/examples/simple_context_bandit.py", + tags = ["examples", "examples_U"], + size = "small", + srcs = ["contrib/bandits/examples/simple_context_bandit.py"], + args = ["--stop-at-reward=10", "--run=contrib/LinUCB"], +) + py_test( name = "examples/twostep_game_pg", main = "examples/twostep_game.py", tags = ["examples", "examples_T"], diff --git a/rllib/contrib/bandits/__init__.py b/rllib/contrib/bandits/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/contrib/bandits/agents/__init__.py b/rllib/contrib/bandits/agents/__init__.py new file mode 100644 index 000000000..70a26104f --- /dev/null +++ b/rllib/contrib/bandits/agents/__init__.py @@ -0,0 +1,4 @@ +from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer +from ray.rllib.contrib.bandits.agents.lin_ucb import LinUCBTrainer + +__all__ = ["LinTSTrainer", "LinUCBTrainer"] diff --git a/rllib/contrib/bandits/agents/lin_ts.py b/rllib/contrib/bandits/agents/lin_ts.py new file mode 100644 index 000000000..e237f9209 --- /dev/null +++ b/rllib/contrib/bandits/agents/lin_ts.py @@ -0,0 +1,46 @@ +import logging + +from ray.rllib.agents.trainer import with_common_config +from ray.rllib.agents.trainer_template import build_trainer +from ray.rllib.contrib.bandits.agents.policy import BanditPolicy + +logger = logging.getLogger(__name__) + +# yapf: disable +# __sphinx_doc_begin__ +TS_CONFIG = with_common_config({ + # No remote workers by default. + "num_workers": 0, + "use_pytorch": True, + + # Do online learning one step at a time. + "rollout_fragment_length": 1, + "train_batch_size": 1, + + # Bandits cant afford to do one timestep per iteration as it is extremely + # slow because of metrics collection overhead. This setting means that the + # agent will be trained for 100 times in one iteration of Rllib + "timesteps_per_iteration": 100, + + "exploration_config": { + "type": "ray.rllib.contrib.bandits.exploration.ThompsonSampling" + } +}) + +# __sphinx_doc_end__ +# yapf: enable + + +def get_stats(trainer): + env_metrics = trainer.collect_metrics() + stats = trainer.optimizer.stats() + # Uncomment if regret at each time step is needed + # stats.update({"all_regrets": trainer.get_policy().regrets}) + return dict(env_metrics, **stats) + + +LinTSTrainer = build_trainer( + name="LinTS", + default_config=TS_CONFIG, + default_policy=BanditPolicy, + collect_metrics_fn=get_stats) diff --git a/rllib/contrib/bandits/agents/lin_ucb.py b/rllib/contrib/bandits/agents/lin_ucb.py new file mode 100644 index 000000000..36029c3fa --- /dev/null +++ b/rllib/contrib/bandits/agents/lin_ucb.py @@ -0,0 +1,46 @@ +import logging + +from ray.rllib.agents.trainer import with_common_config +from ray.rllib.agents.trainer_template import build_trainer +from ray.rllib.contrib.bandits.agents.policy import BanditPolicy + +logger = logging.getLogger(__name__) + +# yapf: disable +# __sphinx_doc_begin__ +UCB_CONFIG = with_common_config({ + # No remote workers by default. + "num_workers": 0, + "use_pytorch": True, + + # Do online learning one step at a time. + "rollout_fragment_length": 1, + "train_batch_size": 1, + + # Bandits cant afford to do one timestep per iteration as it is extremely + # slow because of metrics collection overhead. This setting means that the + # agent will be trained for 100 times in one iteration of Rllib + "timesteps_per_iteration": 100, + + "exploration_config": { + "type": "ray.rllib.contrib.bandits.exploration.UCB" + } +}) + +# __sphinx_doc_end__ +# yapf: enable + + +def get_stats(trainer): + env_metrics = trainer.collect_metrics() + stats = trainer.optimizer.stats() + # Uncomment if regret at each time step is needed + # stats.update({"all_regrets": trainer.get_policy().regrets}) + return dict(env_metrics, **stats) + + +LinUCBTrainer = build_trainer( + name="LinUCB", + default_config=UCB_CONFIG, + default_policy=BanditPolicy, + collect_metrics_fn=get_stats) diff --git a/rllib/contrib/bandits/agents/policy.py b/rllib/contrib/bandits/agents/policy.py new file mode 100644 index 000000000..d0144a466 --- /dev/null +++ b/rllib/contrib/bandits/agents/policy.py @@ -0,0 +1,121 @@ +import logging +import time + +from gym import spaces +from ray.rllib.agents.trainer import with_common_config +from ray.rllib.contrib.bandits.models.linear_regression import \ + DiscreteLinearModelThompsonSampling, \ + DiscreteLinearModelUCB, DiscreteLinearModel, \ + ParametricLinearModelThompsonSampling, ParametricLinearModelUCB +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.model import restore_original_dimensions +from ray.rllib.policy.policy import LEARNER_STATS_KEY +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_policy import TorchPolicy +from ray.rllib.policy.torch_policy_template import build_torch_policy +from ray.rllib.utils.annotations import override +from ray.util.debug import log_once + +logger = logging.getLogger(__name__) + + +TS_PATH = "ray.rllib.contrib.bandits.exploration.ThompsonSampling" +UCB_PATH = "ray.rllib.contrib.bandits.exploration.UCB" + + +DEFAULT_CONFIG = with_common_config({ + # No remote workers by default. + "num_workers": 0, + "use_pytorch": True, + + # Do online learning one step at a time. + "rollout_fragment_length": 1, + "train_batch_size": 1, + + # Bandits cant afford to do one timestep per iteration as it is extremely + # slow because of metrics collection overhead. This setting means that the + # agent will be trained for 100 times in one iteration of Rllib + "timesteps_per_iteration": 100 +}) + + +class BanditPolicyOverrides: + @override(TorchPolicy) + def learn_on_batch(self, postprocessed_batch): + train_batch = self._lazy_tensor_dict(postprocessed_batch) + unflattened_obs = restore_original_dimensions( + train_batch[SampleBatch.CUR_OBS], self.observation_space, + self.framework) + + info = {} + + start = time.time() + self.model.partial_fit(unflattened_obs, + train_batch[SampleBatch.REWARDS], + train_batch[SampleBatch.ACTIONS]) + + infos = postprocessed_batch["infos"] + if "regret" in infos[0]: + regret = sum( + row["infos"]["regret"] for row in postprocessed_batch.rows()) + self.regrets.append(regret) + info["cumulative_regret"] = sum(self.regrets) + else: + if log_once("no_regrets"): + logger.warning("The env did not report `regret` values in " + "its `info` return, ignoring.") + info["update_latency"] = time.time() - start + return {LEARNER_STATS_KEY: info} + + +def make_model_and_action_dist(policy, obs_space, action_space, config): + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, config["model"], framework="torch") + model_cls = DiscreteLinearModel + + if hasattr(obs_space, "original_space"): + original_space = obs_space.original_space + else: + original_space = obs_space + + exploration_config = config.get("exploration_config") + # Model is dependent on exploration strategy because of its implicitness + + # TODO: Have a separate model catalogue for bandits + if exploration_config: + if exploration_config["type"] == TS_PATH: + if isinstance(original_space, spaces.Dict): + assert "item" in original_space.spaces, \ + "Cannot find 'item' key in observation space" + model_cls = ParametricLinearModelThompsonSampling + else: + model_cls = DiscreteLinearModelThompsonSampling + elif exploration_config["type"] == UCB_PATH: + if isinstance(original_space, spaces.Dict): + assert "item" in original_space.spaces, \ + "Cannot find 'item' key in observation space" + model_cls = ParametricLinearModelUCB + else: + model_cls = DiscreteLinearModelUCB + + model = model_cls( + obs_space, + action_space, + logit_dim, + config["model"], + name="LinearModel") + return model, dist_class + + +def init_cum_regret(policy, *args): + policy.regrets = [] + + +BanditPolicy = build_torch_policy( + name="BanditPolicy", + get_default_config=lambda: DEFAULT_CONFIG, + loss_fn=None, + after_init=init_cum_regret, + make_model_and_action_dist=make_model_and_action_dist, + optimizer_fn=lambda policy, config: None, # Pass a dummy optimizer + mixins=[BanditPolicyOverrides]) diff --git a/rllib/contrib/bandits/envs/__init__.py b/rllib/contrib/bandits/envs/__init__.py new file mode 100644 index 000000000..9ce1c5f05 --- /dev/null +++ b/rllib/contrib/bandits/envs/__init__.py @@ -0,0 +1,5 @@ +from ray.rllib.contrib.bandits.envs.discrete import LinearDiscreteEnv, \ + WheelBanditEnv +from ray.rllib.contrib.bandits.envs.parametric import ParametricItemRecoEnv + +__all__ = ["LinearDiscreteEnv", "WheelBanditEnv", "ParametricItemRecoEnv"] diff --git a/rllib/contrib/bandits/envs/discrete.py b/rllib/contrib/bandits/envs/discrete.py new file mode 100644 index 000000000..b0f2f4277 --- /dev/null +++ b/rllib/contrib/bandits/envs/discrete.py @@ -0,0 +1,171 @@ +import copy + +import gym +import numpy as np +from gym import spaces + +DEFAULT_CONFIG_LINEAR = { + "feature_dim": 8, + "num_actions": 4, + "reward_noise_std": 0.01 +} + + +class LinearDiscreteEnv(gym.Env): + """Samples data from linearly parameterized arms. + + The reward for context X and arm i is given by X^T * theta_i, for some + latent set of parameters {theta_i : i = 1, ..., k}. The betas are sampled + uniformly at random, the contexts are Gaussian, and Gaussian noise is + added to the rewards. + + """ + + def __init__(self, config=None): + self.config = copy.copy(DEFAULT_CONFIG_LINEAR) + if config is not None and type(config) == dict: + self.config.update(config) + + self.feature_dim = self.config["feature_dim"] + self.num_actions = self.config["num_actions"] + self.sigma = self.config["reward_noise_std"] + + self.action_space = spaces.Discrete(self.num_actions) + self.observation_space = spaces.Box( + low=-10, high=10, shape=(self.feature_dim, )) + + self.thetas = np.random.uniform(-1, 1, + (self.num_actions, self.feature_dim)) + self.thetas /= np.linalg.norm(self.thetas, axis=1, keepdims=True) + + self._elapsed_steps = 0 + self._current_context = None + + def _sample_context(self): + return np.random.normal(scale=1 / 3, size=(self.feature_dim, )) + + def reset(self): + self._current_context = self._sample_context() + return self._current_context + + def step(self, action): + assert self._elapsed_steps is not None,\ + "Cannot call env.step() beforecalling reset()" + assert action < self.num_actions, "Invalid action." + + action = int(action) + context = self._current_context + rewards = self.thetas.dot(context) + + opt_action = rewards.argmax() + + regret = rewards.max() - rewards[action] + + # Add Gaussian noise + rewards += np.random.normal(scale=self.sigma, size=rewards.shape) + + reward = rewards[action] + self._current_context = self._sample_context() + return self._current_context, reward, True, { + "regret": regret, + "opt_action": opt_action + } + + def render(self, mode="human"): + raise NotImplementedError + + +DEFAULT_CONFIG_WHEEL = { + "delta": 0.5, + "mu_1": 1.2, + "mu_2": 1, + "mu_3": 50, + "std": 0.01 +} + + +class WheelBanditEnv(gym.Env): + """Wheel bandit environment for 2D contexts + (see https://arxiv.org/abs/1802.09127). + """ + + feature_dim = 2 + num_actions = 5 + + def __init__(self, config=None): + self.config = copy.copy(DEFAULT_CONFIG_WHEEL) + if config is not None and type(config) == dict: + self.config.update(config) + + self.delta = self.config["delta"] + self.mu_1 = self.config["mu_1"] + self.mu_2 = self.config["mu_2"] + self.mu_3 = self.config["mu_3"] + self.std = self.config["std"] + + self.action_space = spaces.Discrete(self.num_actions) + self.observation_space = spaces.Box( + low=-1, high=1, shape=(self.feature_dim, )) + + self.means = [self.mu_1] + 4 * [self.mu_2] + self._elapsed_steps = 0 + self._current_context = None + + def _sample_context(self): + while True: + state = np.random.uniform(-1, 1, self.feature_dim) + if np.linalg.norm(state) <= 1: + return state + + def reset(self): + self._current_context = self._sample_context() + return self._current_context + + def step(self, action): + assert self._elapsed_steps is not None,\ + "Cannot call env.step() before calling reset()" + + action = int(action) + self._elapsed_steps += 1 + rewards = [ + np.random.normal(self.means[j], self.std) + for j in range(self.num_actions) + ] + context = self._current_context + r_big = np.random.normal(self.mu_3, self.std) + + if np.linalg.norm(context) >= self.delta: + if context[0] > 0: + if context[1] > 0: + # First quadrant + rewards[1] = r_big + opt_action = 1 + else: + # Fourth quadrant + rewards[4] = r_big + opt_action = 4 + else: + if context[1] > 0: + # Second quadrant + rewards[2] = r_big + opt_action = 2 + else: + # Third quadrant + rewards[3] = r_big + opt_action = 3 + else: + # Smaller region where action 0 is optimal + opt_action = 0 + + reward = rewards[action] + + regret = rewards[opt_action] - reward + + self._current_context = self._sample_context() + return self._current_context, reward, True, { + "regret": regret, + "opt_action": opt_action + } + + def render(self, mode="human"): + raise NotImplementedError diff --git a/rllib/contrib/bandits/envs/parametric.py b/rllib/contrib/bandits/envs/parametric.py new file mode 100644 index 000000000..a187d9134 --- /dev/null +++ b/rllib/contrib/bandits/envs/parametric.py @@ -0,0 +1,157 @@ +import copy + +import gym +import numpy as np +from gym import spaces + +DEFAULT_RECO_CONFIG = { + "num_users": 1, + "num_items": 100, + "feature_dim": 16, + "slate_size": 1, + "num_candidates": 25, + "seed": 1 +} + + +class ParametricItemRecoEnv(gym.Env): + """A recommendation environment which generates items with visible features + randomly (parametric actions). + The environment can be configured to be multi-user, i.e. different model + will be learned independently for each user. + To enable slate recommendation, the `slate_size` config parameter can be + set as > 1 . + """ + + def __init__(self, config=None): + self.config = copy.copy(DEFAULT_RECO_CONFIG) + if config is not None and type(config) == dict: + self.config.update(config) + + self.num_users = self.config["num_users"] + self.num_items = self.config["num_items"] + self.feature_dim = self.config["feature_dim"] + self.slate_size = self.config["slate_size"] + self.num_candidates = self.config["num_candidates"] + self.seed = self.config["seed"] + + assert self.num_candidates <= self.num_items,\ + "Size of candidate pool should be less than total no. of items" + assert self.slate_size < self.num_candidates,\ + "Slate size should be less than no. of candidate items" + + self.action_space = self._def_action_space() + self.observation_space = self._def_observation_space() + + self.current_user_id = 0 + self.item_pool = None + self.item_pool_ids = None + self.total_regret = 0 + + self._init_embeddings() + + def _init_embeddings(self): + self.item_embeddings = self._gen_normalized_embeddings( + self.num_items, self.feature_dim) + + # These are latent user features that will be hidden from the learning + # agent. They will be used for reward generation only + self.user_embeddings = self._gen_normalized_embeddings( + self.num_users, self.feature_dim) + + def _sample_user(self): + self.current_user_id = np.random.randint(0, self.num_users) + + def _gen_item_pool(self): + # Randomly generate a candidate list of items by sampling without + # replacement + self.item_pool_ids = np.random.choice( + np.arange(self.num_items), self.num_candidates, replace=False) + self.item_pool = self.item_embeddings[self.item_pool_ids] + + @staticmethod + def _gen_normalized_embeddings(size, dim): + embeddings = np.random.rand(size, dim) + embeddings /= np.linalg.norm(embeddings, axis=1, keepdims=True) + return embeddings + + def _def_action_space(self): + if self.slate_size == 1: + return spaces.Discrete(self.num_candidates) + else: + return spaces.MultiDiscrete( + [self.num_candidates] * self.slate_size) + + def _def_observation_space(self): + # Embeddings for each item in the candidate pool + item_obs_space = spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.num_candidates, self.feature_dim)) + + # Can be useful for collaborative filtering based agents + item_ids_obs_space = spaces.MultiDiscrete( + [self.num_items] * self.num_candidates) + + # Can be either binary (clicks) or continuous feedback (watch time) + resp_space = spaces.Box(low=-1, high=1, shape=(self.slate_size, )) + + if self.num_users == 1: + return spaces.Dict({ + "item": item_obs_space, + "item_id": item_ids_obs_space, + "response": resp_space + }) + else: + user_obs_space = spaces.Discrete(self.num_users) + return spaces.Dict({ + "user": user_obs_space, + "item": item_obs_space, + "item_id": item_ids_obs_space, + "response": resp_space + }) + + def step(self, action): + # Action can be a single action or a slate depending on slate size + assert self.action_space.contains( + action + ), "Action cannot be recognized. Please check the type and bounds." + + if self.slate_size == 1: + scores = self.item_pool.dot( + self.user_embeddings[self.current_user_id]) + reward = scores[action] + regret = np.max(scores) - reward + self.total_regret += regret + + info = {"regret": regret} + + self.current_user_id = np.random.randint(0, self.num_users) + self._gen_item_pool() + + obs = { + "item": self.item_pool, + "item_id": self.item_pool_ids, + "response": [reward] + } + if self.num_users > 1: + obs["user"] = self.current_user_id + return obs, reward, True, info + else: + # TODO(saurabh3949):Handle slate recommendation using a click model + return None + + def reset(self): + self._sample_user() + self._gen_item_pool() + obs = { + "item": self.item_pool, + "item_id": self.item_pool_ids, + "response": [0] * self.slate_size + } + if self.num_users > 1: + obs["user"] = self.current_user_id + return obs + + def render(self, mode="human"): + raise NotImplementedError diff --git a/rllib/contrib/bandits/examples/LinTS_train_wheel_env.py b/rllib/contrib/bandits/examples/LinTS_train_wheel_env.py new file mode 100644 index 000000000..d9ab5315a --- /dev/null +++ b/rllib/contrib/bandits/examples/LinTS_train_wheel_env.py @@ -0,0 +1,52 @@ +""" Example of using Linear Thompson Sampling on WheelBandit environment. + For more information on WheelBandit, see https://arxiv.org/abs/1802.09127 . +""" + +import numpy as np +from matplotlib import pyplot as plt +from ray.rllib.contrib.bandits.agents import LinTSTrainer +from ray.rllib.contrib.bandits.envs import WheelBanditEnv + + +def plot_model_weights(means, covs): + fmts = ["bo", "ro", "yx", "k+", "gx"] + labels = [f"arm{i}" for i in range(5)] + + fig, ax = plt.subplots(figsize=(6, 4)) + + ax.set_title("Weights distributions of arms") + + for i in range(0, 5): + x, y = np.random.multivariate_normal(means[i] / 30, covs[i], 5000).T + ax.plot(x, y, fmts[i], label=labels[i]) + + ax.grid(True, which="both") + ax.axhline(y=0, color="k") + ax.axvline(x=0, color="k") + ax.legend(loc="best") + plt.show() + + +if __name__ == "__main__": + num_iter = 20 + print("Running training for %s time steps" % num_iter) + trainer = LinTSTrainer(env=WheelBanditEnv) + + policy = trainer.get_policy() + model = policy.model + + print("Using exploration strategy:", policy.exploration) + print("Using model:", model) + + for i in range(num_iter): + trainer.train() + + info = trainer.train() + print(info["learner"]) + + # Get model parameters + means = [model.arms[i].theta.numpy() for i in range(5)] + covs = [model.arms[i].covariance.numpy() for i in range(5)] + + # Plot weight distributions for different arms + plot_model_weights(means, covs) diff --git a/rllib/contrib/bandits/examples/simple_context_bandit.py b/rllib/contrib/bandits/examples/simple_context_bandit.py new file mode 100644 index 000000000..9ea5357f9 --- /dev/null +++ b/rllib/contrib/bandits/examples/simple_context_bandit.py @@ -0,0 +1,47 @@ +"""A very simple contextual bandit example with 3 arms.""" + +import argparse +import random +import numpy as np +import gym +from gym.spaces import Discrete, Box + +from ray import tune + +parser = argparse.ArgumentParser() +parser.add_argument("--stop-at-reward", type=float, default=10) +parser.add_argument("--run", type=str, default="contrib/LinUCB") + + +class SimpleContextualBandit(gym.Env): + def __init__(self, config=None): + self.action_space = Discrete(3) + self.observation_space = Box(low=-1., high=1., shape=(2, )) + self.cur_context = None + + def reset(self): + self.cur_context = random.choice([-1., 1.]) + return np.array([self.cur_context, -self.cur_context]) + + def step(self, action): + rewards_for_context = { + -1.: [-10, 0, 10], + 1.: [10, 0, -10], + } + reward = rewards_for_context[self.cur_context][action] + return (np.array([-self.cur_context, self.cur_context]), reward, True, + { + "regret": 10 - reward + }) + + +if __name__ == "__main__": + args = parser.parse_args() + tune.run( + args.run, + stop={ + "episode_reward_mean": args.stop_at_reward, + }, + config={ + "env": SimpleContextualBandit, + }) diff --git a/rllib/contrib/bandits/examples/tune_LinTS_train_wheel_env.py b/rllib/contrib/bandits/examples/tune_LinTS_train_wheel_env.py new file mode 100644 index 000000000..b319ba444 --- /dev/null +++ b/rllib/contrib/bandits/examples/tune_LinTS_train_wheel_env.py @@ -0,0 +1,80 @@ +""" Example of using Linear Thompson Sampling on WheelBandit environment. + For more information on WheelBandit, see https://arxiv.org/abs/1802.09127 . +""" + +import time + +import numpy as np +import pandas as pd +from matplotlib import pyplot as plt +from ray import tune +from ray.rllib.contrib.bandits.agents import LinTSTrainer +from ray.rllib.contrib.bandits.agents.lin_ts import TS_CONFIG +from ray.rllib.contrib.bandits.envs import WheelBanditEnv + + +def plot_model_weights(means, covs, ax): + fmts = ["bo", "ro", "yx", "k+", "gx"] + labels = [f"arm{i}" for i in range(5)] + + ax.set_title("Weights distributions of arms") + + for i in range(0, 5): + x, y = np.random.multivariate_normal(means[i] / 30, covs[i], 5000).T + ax.plot(x, y, fmts[i], label=labels[i]) + + ax.set_aspect("equal") + ax.grid(True, which="both") + ax.axhline(y=0, color="k") + ax.axvline(x=0, color="k") + ax.legend(loc="best") + + +if __name__ == "__main__": + TS_CONFIG["env"] = WheelBanditEnv + + # Actual training_iterations will be 20 * timesteps_per_iteration + # (100 by default) = 2,000 + training_iterations = 20 + + print("Running training for %s time steps" % training_iterations) + + start_time = time.time() + analysis = tune.run( + LinTSTrainer, + config=TS_CONFIG, + stop={"training_iteration": training_iterations}, + num_samples=2, + checkpoint_at_end=True) + + print("The trials took", time.time() - start_time, "seconds\n") + + # Analyze cumulative regrets of the trials + frame = pd.DataFrame() + for key, df in analysis.trial_dataframes.items(): + frame = frame.append(df, ignore_index=True) + + x = frame.groupby("num_steps_trained")[ + "learner/cumulative_regret"].aggregate(["mean", "max", "min", "std"]) + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4)) + + ax1.plot(x["mean"]) + + ax1.set_title("Cumulative Regret") + ax1.set_xlabel("Training steps") + + # Restore trainer from checkpoint + trial = analysis.trials[0] + trainer = LinTSTrainer(config=TS_CONFIG) + trainer.restore(trial.checkpoint.value) + + # Get model to plot arm weights distribution + model = trainer.get_policy().model + means = [model.arms[i].theta.numpy() for i in range(5)] + covs = [model.arms[i].covariance.numpy() for i in range(5)] + + # Plot weight distributions for different arms + plot_model_weights(means, covs, ax2) + fig.tight_layout() + plt.show() diff --git a/rllib/contrib/bandits/examples/tune_LinUCB_train_recommendation.py b/rllib/contrib/bandits/examples/tune_LinUCB_train_recommendation.py new file mode 100644 index 000000000..c27f047bc --- /dev/null +++ b/rllib/contrib/bandits/examples/tune_LinUCB_train_recommendation.py @@ -0,0 +1,54 @@ +""" Example of using LinUCB on a recommendation environment with parametric + actions. """ + +import os +import time + +from matplotlib import pyplot as plt +import pandas as pd + +from ray import tune +from ray.rllib.contrib.bandits.agents import LinUCBTrainer +from ray.rllib.contrib.bandits.agents.lin_ucb import UCB_CONFIG +from ray.rllib.contrib.bandits.envs import ParametricItemRecoEnv + +if __name__ == "__main__": + + # Temp fix to avoid OMP conflict + os.environ["KMP_DUPLICATE_LIB_OK"] = "True" + + UCB_CONFIG["env"] = ParametricItemRecoEnv + + # Actual training_iterations will be 20 * timesteps_per_iteration + # (100 by default) = 2,000 + training_iterations = 20 + + print("Running training for %s time steps" % training_iterations) + + start_time = time.time() + analysis = tune.run( + LinUCBTrainer, + config=UCB_CONFIG, + stop={"training_iteration": training_iterations}, + num_samples=5, + checkpoint_at_end=False) + + print("The trials took", time.time() - start_time, "seconds\n") + + # Analyze cumulative regrets of the trials + frame = pd.DataFrame() + for key, df in analysis.trial_dataframes.items(): + frame = frame.append(df, ignore_index=True) + x = frame.groupby("num_steps_trained")[ + "learner/cumulative_regret"].aggregate(["mean", "max", "min", "std"]) + + plt.plot(x["mean"]) + plt.fill_between( + x.index, + x["mean"] - x["std"], + x["mean"] + x["std"], + color="b", + alpha=0.2) + plt.title("Cumulative Regret") + plt.xlabel("Training steps") + plt.show() diff --git a/rllib/contrib/bandits/exploration.py b/rllib/contrib/bandits/exploration.py new file mode 100644 index 000000000..099b56a03 --- /dev/null +++ b/rllib/contrib/bandits/exploration.py @@ -0,0 +1,52 @@ +from typing import Union + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import TensorType + + +class ThompsonSampling(Exploration): + @override(Exploration) + def get_exploration_action(self, + distribution_inputs: TensorType, + action_dist_class: type, + model: ModelV2, + timestep: Union[int, TensorType], + explore: bool = True): + if self.framework == "torch": + return self._get_torch_exploration_action(distribution_inputs, + explore, model) + else: + raise NotImplementedError + + def _get_torch_exploration_action(self, distribution_inputs, explore, + model): + if explore: + return distribution_inputs.argmax(dim=1), None + else: + scores = model.predict(model.current_obs()) + return scores.argmax(dim=1), None + + +class UCB(Exploration): + @override(Exploration) + def get_exploration_action(self, + distribution_inputs: TensorType, + action_dist_class: type, + model: ModelV2, + timestep: Union[int, TensorType], + explore: bool = True): + if self.framework == "torch": + return self._get_torch_exploration_action(distribution_inputs, + explore, model) + else: + raise NotImplementedError + + def _get_torch_exploration_action(self, distribution_inputs, explore, + model): + if explore: + return distribution_inputs.argmax(dim=1), None + else: + scores = model.value_function() + return scores.argmax(dim=1), None diff --git a/rllib/contrib/bandits/models/__init__.py b/rllib/contrib/bandits/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/contrib/bandits/models/linear_regression.py b/rllib/contrib/bandits/models/linear_regression.py new file mode 100644 index 000000000..10c0f969f --- /dev/null +++ b/rllib/contrib/bandits/models/linear_regression.py @@ -0,0 +1,246 @@ +import gym +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils import try_import_torch +from ray.rllib.utils.annotations import override + +torch, nn = try_import_torch() + + +class OnlineLinearRegression(nn.Module): + def __init__(self, feature_dim, alpha=1, lambda_=1): + super(OnlineLinearRegression, self).__init__() + self.d = feature_dim + self.alpha = alpha + self.precision = nn.Parameter( + data=lambda_ * torch.eye(self.d), requires_grad=False) + self.f = nn.Parameter(data=torch.zeros(self.d, ), requires_grad=False) + self.covariance = nn.Parameter( + data=torch.inverse(self.precision), requires_grad=False) + self.theta = nn.Parameter( + data=self.covariance.matmul(self.f), requires_grad=False) + self._init_params() + + def _init_params(self): + self.update_schedule = 1 + self.delta_f = 0 + self.delta_b = 0 + self.time = 0 + self.covariance.mul_(self.alpha) + self.dist = torch.distributions.multivariate_normal\ + .MultivariateNormal(self.theta, self.covariance) + + def partial_fit(self, x, y): + # TODO: Handle batch of data rather than individual points + self._check_inputs(x, y) + x = x.squeeze() + y = y.item() + self.time += 1 + self.delta_f += y * x + self.delta_b += torch.ger(x, x) + # Can follow an update schedule if not doing sherman morison updates + if self.time % self.update_schedule == 0: + self.precision += self.delta_b + self.f += self.delta_f + self.delta_b = 0 + self.delta_f = 0 + torch.inverse(self.precision, out=self.covariance) + torch.matmul(self.covariance, self.f, out=self.theta) + self.covariance.mul_(self.alpha) + + def sample_theta(self): + theta = self.dist.sample() + return theta + + def get_ucbs(self, x): + """ Calculate upper confidence bounds using covariance matrix according + to algorithm 1: LinUCB + (http://proceedings.mlr.press/v15/chu11a/chu11a.pdf). + + Args: + x (torch.Tensor): Input feature tensor of shape + (batch_size, feature_dim) + """ + + projections = self.covariance @ x.T + batch_dots = (x * projections.T).sum(dim=1) + return batch_dots.sqrt() + + def forward(self, x, sample_theta=False): + """ Predict the scores on input batch using the underlying linear model + Args: + x (torch.Tensor): Input feature tensor of shape + (batch_size, feature_dim) + sample_theta (bool): Whether to sample the weights from its + posterior distribution to perform Thompson Sampling as per + http://proceedings.mlr.press/v28/agrawal13.pdf . + """ + self._check_inputs(x) + theta = self.sample_theta() if sample_theta else self.theta + scores = x @ theta + return scores + + def _check_inputs(self, x, y=None): + assert x.ndim in [2, 3], \ + "Input context tensor must be 2 or 3 dimensional, where the" \ + " first dimension is batch size" + assert x.shape[ + 1] == self.d, f"Feature dimensions of weights ({self.d}) and " \ + f"context ({x.shape[1]}) do not match!" + if y: + assert torch.is_tensor(y) and y.numel() == 1,\ + "Target should be a tensor;" \ + "Only online learning with a batch size of 1 is " \ + "supported for now!" + + +class DiscreteLinearModel(TorchModelV2, nn.Module): + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + + alpha = model_config.get("alpha", 1) + lambda_ = model_config.get("lambda_", 1) + self.feature_dim = obs_space.sample().size + self.arms = nn.ModuleList([ + OnlineLinearRegression( + feature_dim=self.feature_dim, alpha=alpha, lambda_=lambda_) + for i in range(self.num_outputs) + ]) + self._cur_value = None + self._cur_ctx = None + + @override(TorchModelV2) + def forward(self, input_dict, state, seq_lens): + x = input_dict["obs"] + scores = self.predict(x) + return scores, state + + def predict(self, x, sample_theta=False, use_ucb=False): + self._cur_ctx = x + scores = torch.stack( + [self.arms[i](x, sample_theta) for i in range(self.num_outputs)], + dim=-1) + self._cur_value = scores + if use_ucb: + ucbs = torch.stack( + [self.arms[i].get_ucbs(x) for i in range(self.num_outputs)], + dim=-1) + return scores + ucbs + else: + return scores + + def partial_fit(self, x, y, arm): + assert 0 <= arm.item() < len(self.arms),\ + f"Invalid arm: {arm.item()}." \ + f"It should be 0 <= arm < {len(self.arms)}" + self.arms[arm].partial_fit(x, y) + + @override(TorchModelV2) + def value_function(self): + assert self._cur_value is not None, "must call forward() first" + return self._cur_value + + def current_obs(self): + assert self._cur_ctx is not None, "must call forward() first" + return self._cur_ctx + + +class DiscreteLinearModelUCB(DiscreteLinearModel): + def forward(self, input_dict, state, seq_lens): + x = input_dict["obs"] + scores = super(DiscreteLinearModelUCB, self).predict( + x, sample_theta=False, use_ucb=True) + return scores, state + + +class DiscreteLinearModelThompsonSampling(DiscreteLinearModel): + def forward(self, input_dict, state, seq_lens): + x = input_dict["obs"] + scores = super(DiscreteLinearModelThompsonSampling, self).predict( + x, sample_theta=True, use_ucb=False) + return scores, state + + +class ParametricLinearModel(TorchModelV2, nn.Module): + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + + alpha = model_config.get("alpha", 1) + lambda_ = model_config.get("lambda_", 0.1) + + # RLlib preprocessors will flatten the observation space and unflatten + # it later. Accessing the original space here. + original_space = obs_space.original_space + assert isinstance(original_space, gym.spaces.Dict) and \ + "item" in original_space.spaces, \ + "This model only supports gym.spaces.Dict observation spaces." + self.feature_dim = original_space["item"].shape[-1] + self.arm = OnlineLinearRegression( + feature_dim=self.feature_dim, alpha=alpha, lambda_=lambda_) + self._cur_value = None + self._cur_ctx = None + + def _check_inputs(self, x): + if x.ndim == 3: + assert x.size()[ + 0] == 1, "Only batch size of 1 is supported for now." + + @override(TorchModelV2) + def forward(self, input_dict, state, seq_lens): + x = input_dict["obs"]["item"] + self._check_inputs(x) + x.squeeze_(dim=0) # Remove the batch dimension + scores = self.predict(x) + scores.unsqueeze_(dim=0) # Add the batch dimension + return scores, state + + def predict(self, x, sample_theta=False, use_ucb=False): + self._cur_ctx = x + scores = self.arm(x, sample_theta) + self._cur_value = scores + if use_ucb: + ucbs = self.arm.get_ucbs(x) + return scores + 0.3 * ucbs + else: + return scores + + def partial_fit(self, x, y, arm): + x = x["item"] + action_id = arm.item() + self.arm.partial_fit(x[:, action_id], y) + + @override(TorchModelV2) + def value_function(self): + assert self._cur_value is not None, "must call forward() first" + return self._cur_value + + def current_obs(self): + assert self._cur_ctx is not None, "must call forward() first" + return self._cur_ctx + + +class ParametricLinearModelUCB(ParametricLinearModel): + def forward(self, input_dict, state, seq_lens): + x = input_dict["obs"]["item"] + self._check_inputs(x) + x.squeeze_(dim=0) # Remove the batch dimension + scores = super(ParametricLinearModelUCB, self).predict( + x, sample_theta=False, use_ucb=True) + scores.unsqueeze_(dim=0) # Add the batch dimension + return scores, state + + +class ParametricLinearModelThompsonSampling(ParametricLinearModel): + def forward(self, input_dict, state, seq_lens): + x = input_dict["obs"]["item"] + self._check_inputs(x) + x.squeeze_(dim=0) # Remove the batch dimension + scores = super(ParametricLinearModelThompsonSampling, self).predict( + x, sample_theta=True, use_ucb=False) + scores.unsqueeze_(dim=0) # Add the batch dimension + return scores, state diff --git a/rllib/contrib/registry.py b/rllib/contrib/registry.py index f2c0fca49..aed8712bb 100644 --- a/rllib/contrib/registry.py +++ b/rllib/contrib/registry.py @@ -17,8 +17,20 @@ def _import_alphazero(): return AlphaZeroTrainer +def _import_bandit_lints(): + from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer + return LinTSTrainer + + +def _import_bandit_linucb(): + from ray.rllib.contrib.bandits.agents.lin_ucb import LinUCBTrainer + return LinUCBTrainer + + CONTRIBUTED_ALGORITHMS = { "contrib/RandomAgent": _import_random_agent, "contrib/MADDPG": _import_maddpg, "contrib/AlphaZero": _import_alphazero, + "contrib/LinTS": _import_bandit_lints, + "contrib/LinUCB": _import_bandit_linucb }