Contextual Bandit algorithms (WIP) (#7642)

.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

doc/source/rllib-algorithms.rst

@@ -465,8 +465,69 @@ Tuned examples: `CartPole-v0 <https://github.com/ray-project/ray/blob/master/rll
    :start-after: __sphinx_doc_begin__
    :end-before: __sphinx_doc_end__
 
+
+Contextual Bandits (contrib/bandits)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The Multi-armed bandit (MAB) problem provides a simplified RL setting that
+involves learning to act under one situation only, i.e. the state is fixed.
+Contextual bandit is extension of the MAB problem, where at each
+round the agent has access not only to a set of bandit arms/actions but also
+to a context (state) associated with this iteration. The context changes
+with each iteration, but, is not affected by the action that the agent takes.
+The objective of the agent is to maximize the cumulative rewards, by
+collecting  enough information about how the context and the rewards of the
+arms are related to each other. The agent does this by balancing the
+trade-off between exploration and exploitation.
+
+Contextual bandit algorithms typically consist of an action-value model (Q
+model) and an exploration strategy (e-greedy, UCB, Thompson Sampling etc.)
+RLlib supports the following online contextual bandit algorithms,
+named after the exploration strategies that they employ:
+
+LinUCB (Upper Confidence Bound)
+-------------------------------
+|pytorch|
+`[paper] <http://rob.schapire.net/papers/www10.pdf>`__ `[implementation]
+<https://github.com/ray-project/ray/blob/master/rllib/contrib/bandits/agents/lin_ucb.py>`__
+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 <rllib-training
+.html#common-parameters>`__):
+
+.. 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] <http://proceedings.mlr.press/v28/agrawal13.pdf>`__ `[implementation]
+<https://github.com/ray-project/ray/blob/master/rllib/contrib/bandits/agents/lin_ts.py>`__
+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 <rllib-training
+.html#common-parameters>`__):
+
+.. 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

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

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

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)

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)

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])

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

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

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

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)

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,
+        })

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()

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()

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

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

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
 }