diff --git a/CHANGELOG.txt b/CHANGELOG.txt
index f6e13a1..fbe2404 100755
--- a/CHANGELOG.txt
+++ b/CHANGELOG.txt
@@ -2,6 +2,14 @@
 MABWiser CHANGELOG
 =====================
 
+March, 17, 2022 2.4.0
+-------------------------------------------------------------------------------
+major:
+- Implement vectorized functions for non-contextual policies to speed-up prediction for multiple decisions.
+- Change MAB predict and predict_expectations to allow empty contexts to be specified for non-contextual policies.
+- Update scaler use in Linear policies so that standard scaler can be fit directly instead of pre-trained scalers.
+- Change scaler argument from pre-trained `arm_to_scaler` input to a boolean scale flag.
+
 March, 8, 2022 2.3.0
 -------------------------------------------------------------------------------
 major:
diff --git a/docs/.buildinfo b/docs/.buildinfo
index b73da4e..d9e4d7f 100644
--- a/docs/.buildinfo
+++ b/docs/.buildinfo
@@ -1,4 +1,4 @@
 # Sphinx build info version 1
 # This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done.
-config: be8d7ae06467e3a22a202d69cb2b24ab
+config: 3d55ee819e2d2bd60e727943c52c4b22
 tags: 645f666f9bcd5a90fca523b33c5a78b7
diff --git a/docs/_static/documentation_options.js b/docs/_static/documentation_options.js
index 32053ba..2e89cf7 100644
--- a/docs/_static/documentation_options.js
+++ b/docs/_static/documentation_options.js
@@ -1,6 +1,6 @@
 var DOCUMENTATION_OPTIONS = {
     URL_ROOT: document.getElementById("documentation_options").getAttribute('data-url_root'),
-    VERSION: '2.3.0',
+    VERSION: '2.4.0',
     LANGUAGE: 'None',
     COLLAPSE_INDEX: false,
     BUILDER: 'html',
diff --git a/docs/about.html b/docs/about.html
index 0a44bdf..7de1105 100644
--- a/docs/about.html
+++ b/docs/about.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>About Multi-Armed Bandits &mdash; MABWiser 2.3.0 documentation</title>
+  <title>About Multi-Armed Bandits &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/api.html b/docs/api.html
index 24a6b8e..3f85490 100644
--- a/docs/api.html
+++ b/docs/api.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>MABWiser Public API &mdash; MABWiser 2.3.0 documentation</title>
+  <title>MABWiser Public API &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
@@ -146,7 +146,7 @@
 <dt class="sig sig-object py" id="mabwiser.base_mab.BaseMAB.backend">
 <span class="sig-name descname"><span class="pre">backend</span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.backend" title="Permalink to this definition"></a></dt>
 <dd><p>Specify a parallelization backend implementation supported in the joblib library. Supported options are:
-- “loky” used by default, can induce some communication and memory overhead when exchanging input and output data with the worker Python processes.
+- “loky” used by default, can induce some communication and memory overhead when exchanging input and output.
 - “multiprocessing” previous process-based backend based on multiprocessing.Pool. Less robust than loky.
 - “threading” is a very low-overhead backend but it suffers from the Python Global Interpreter Lock if the</p>
 <blockquote>
@@ -196,7 +196,7 @@
 
 <dl class="py method">
 <dt class="sig sig-object py" id="mabwiser.base_mab.BaseMAB.add_arm">
-<span class="sig-name descname"><span class="pre">add_arm</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">arm</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Arm</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">binarizer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">scaler</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">NoReturn</span></span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.add_arm" title="Permalink to this definition"></a></dt>
+<span class="sig-name descname"><span class="pre">add_arm</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">arm</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Arm</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">binarizer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">NoReturn</span></span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.add_arm" title="Permalink to this definition"></a></dt>
 <dd><p>Introduces a new arm to the bandit.</p>
 <p>Adds the new arm with zero expectations and
 calls the <code class="docutils literal notranslate"><span class="pre">_uptake_new_arm()</span></code> function of the sub-class.</p>
@@ -220,14 +220,14 @@
 
 <dl class="py method">
 <dt class="sig sig-object py" id="mabwiser.base_mab.BaseMAB.predict">
-<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">predict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">contexts</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">numpy.ndarray</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Arm</span></span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.predict" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">predict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">contexts</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">numpy.ndarray</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.predict" title="Permalink to this definition"></a></dt>
 <dd><p>Abstract method.</p>
 <p>Returns the predicted arm.</p>
 </dd></dl>
 
 <dl class="py method">
 <dt class="sig sig-object py" id="mabwiser.base_mab.BaseMAB.predict_expectations">
-<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">predict_expectations</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">contexts</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">numpy.ndarray</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.predict_expectations" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">predict_expectations</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">contexts</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">numpy.ndarray</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#mabwiser.base_mab.BaseMAB.predict_expectations" title="Permalink to this definition"></a></dt>
 <dd><p>Abstract method.</p>
 <p>Returns a dictionary from arms (keys) to their expected rewards (values).</p>
 </dd></dl>
@@ -262,7 +262,7 @@
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">tuple</span></code></p>
 <dl class="py class">
 <dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.EpsilonGreedy">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">EpsilonGreedy</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">epsilon</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.05</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.EpsilonGreedy" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">EpsilonGreedy</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">epsilon</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.1</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.EpsilonGreedy" title="Permalink to this definition"></a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">tuple</span></code></p>
 <p>Epsilon Greedy Learning Policy.</p>
 <p>This policy selects the arm with the highest expected reward with probability 1 - <span class="math notranslate nohighlight">\(\epsilon\)</span>,
@@ -272,7 +272,7 @@
 <span class="sig-name descname"><span class="pre">epsilon</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.EpsilonGreedy.epsilon" title="Permalink to this definition"></a></dt>
 <dd><p>The probability of selecting a random arm for exploration.
 Integer or float. Must be between 0 and 1.
-Default value is 0.05.</p>
+Default value is 0.1.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Type</dt>
 <dd class="field-odd"><p>Num</p>
@@ -301,7 +301,7 @@
 
 <dl class="py class">
 <dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinGreedy">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">LinGreedy</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">epsilon</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.05</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">l2_lambda</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">arm_to_scaler</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinGreedy" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">LinGreedy</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">epsilon</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">l2_lambda</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">scale</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinGreedy" title="Permalink to this definition"></a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">tuple</span></code></p>
 <p>LinGreedy Learning Policy.</p>
 <p>This policy trains a ridge regression for each arm.
@@ -313,7 +313,7 @@
 <span class="sig-name descname"><span class="pre">epsilon</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinGreedy.epsilon" title="Permalink to this definition"></a></dt>
 <dd><p>The probability of selecting a random arm for exploration.
 Integer or float. Must be between 0 and 1.
-Default value is 0.05.</p>
+Default value is 0.1.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Type</dt>
 <dd class="field-odd"><p>Num</p>
@@ -335,16 +335,14 @@
 </dd></dl>
 
 <dl class="py attribute">
-<dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinGreedy.arm_to_scaler">
-<span class="sig-name descname"><span class="pre">arm_to_scaler</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinGreedy.arm_to_scaler" title="Permalink to this definition"></a></dt>
-<dd><p>Standardize context features by arm.
-Dictionary mapping each arm to a scaler object. It is assumed
-that the scaler objects are already fit and will only be used
-to transform context features.
-Default value is None.</p>
+<dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinGreedy.scale">
+<span class="sig-name descname"><span class="pre">scale</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinGreedy.scale" title="Permalink to this definition"></a></dt>
+<dd><p>Whether to scale features to have zero mean and unit variance.
+Uses StandardScaler in sklearn.preprocessing.
+Default value is False.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Type</dt>
-<dd class="field-odd"><p>Dict[Arm, Callable]</p>
+<dd class="field-odd"><p>bool</p>
 </dd>
 </dl>
 </dd></dl>
@@ -363,27 +361,27 @@
 </div>
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id1">
-<span class="sig-name descname"><span class="pre">arm_to_scaler</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id1" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 2</p>
+<span class="sig-name descname"><span class="pre">epsilon</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id1" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 0</p>
 </dd></dl>
 
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id2">
-<span class="sig-name descname"><span class="pre">epsilon</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id2" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 0</p>
+<span class="sig-name descname"><span class="pre">l2_lambda</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id2" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 1</p>
 </dd></dl>
 
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id3">
-<span class="sig-name descname"><span class="pre">l2_lambda</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id3" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 1</p>
+<span class="sig-name descname"><span class="pre">scale</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">bool</span></em><a class="headerlink" href="#id3" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 2</p>
 </dd></dl>
 
 </dd></dl>
 
 <dl class="py class">
 <dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinTS">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">LinTS</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">alpha</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">l2_lambda</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">arm_to_scaler</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinTS" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">LinTS</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">alpha</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">l2_lambda</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">scale</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinTS" title="Permalink to this definition"></a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">tuple</span></code></p>
 <p>LinTS Learning Policy</p>
 <p>For each arm LinTS trains a ridge regression and
@@ -426,16 +424,14 @@
 </dd></dl>
 
 <dl class="py attribute">
-<dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinTS.arm_to_scaler">
-<span class="sig-name descname"><span class="pre">arm_to_scaler</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinTS.arm_to_scaler" title="Permalink to this definition"></a></dt>
-<dd><p>Standardize context features by arm.
-Dictionary mapping each arm to a scaler object. It is assumed
-that the scaler objects are already fit and will only be used
-to transform context features.
-Default value is None.</p>
+<dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinTS.scale">
+<span class="sig-name descname"><span class="pre">scale</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinTS.scale" title="Permalink to this definition"></a></dt>
+<dd><p>Whether to scale features to have zero mean and unit variance.
+Uses StandardScaler in sklearn.preprocessing.
+Default value is False.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Type</dt>
-<dd class="field-odd"><p>Dict[Arm, Callable]</p>
+<dd class="field-odd"><p>bool</p>
 </dd>
 </dl>
 </dd></dl>
@@ -460,21 +456,21 @@
 
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id5">
-<span class="sig-name descname"><span class="pre">arm_to_scaler</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id5" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 2</p>
+<span class="sig-name descname"><span class="pre">l2_lambda</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id5" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 1</p>
 </dd></dl>
 
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id6">
-<span class="sig-name descname"><span class="pre">l2_lambda</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id6" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 1</p>
+<span class="sig-name descname"><span class="pre">scale</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">bool</span></em><a class="headerlink" href="#id6" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 2</p>
 </dd></dl>
 
 </dd></dl>
 
 <dl class="py class">
 <dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinUCB">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">LinUCB</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">alpha</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">l2_lambda</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">arm_to_scaler</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinUCB" title="Permalink to this definition"></a></dt>
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">LinUCB</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">alpha</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">l2_lambda</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">scale</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinUCB" title="Permalink to this definition"></a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">tuple</span></code></p>
 <p>LinUCB Learning Policy.</p>
 <p>This policy trains a ridge regression for each arm.
@@ -515,16 +511,14 @@
 </dd></dl>
 
 <dl class="py attribute">
-<dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinUCB.arm_to_scaler">
-<span class="sig-name descname"><span class="pre">arm_to_scaler</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinUCB.arm_to_scaler" title="Permalink to this definition"></a></dt>
-<dd><p>Standardize context features by arm.
-Dictionary mapping each arm to a scaler object. It is assumed
-that the scaler objects are already fit and will only be used
-to transform context features.
-Default value is None.</p>
+<dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.LinUCB.scale">
+<span class="sig-name descname"><span class="pre">scale</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.LinUCB.scale" title="Permalink to this definition"></a></dt>
+<dd><p>Whether to scale features to have zero mean and unit variance.
+Uses StandardScaler in sklearn.preprocessing.
+Default value is False.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Type</dt>
-<dd class="field-odd"><p>Dict[Arm, Callable]</p>
+<dd class="field-odd"><p>bool</p>
 </dd>
 </dl>
 </dd></dl>
@@ -549,14 +543,14 @@
 
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id8">
-<span class="sig-name descname"><span class="pre">arm_to_scaler</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">Arm</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id8" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 2</p>
+<span class="sig-name descname"><span class="pre">l2_lambda</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id8" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 1</p>
 </dd></dl>
 
 <dl class="py attribute">
 <dt class="sig sig-object py" id="id9">
-<span class="sig-name descname"><span class="pre">l2_lambda</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#id9" title="Permalink to this definition"></a></dt>
-<dd><p>Alias for field number 1</p>
+<span class="sig-name descname"><span class="pre">scale</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">bool</span></em><a class="headerlink" href="#id9" title="Permalink to this definition"></a></dt>
+<dd><p>Alias for field number 2</p>
 </dd></dl>
 
 </dd></dl>
@@ -616,8 +610,8 @@
 <div class="math notranslate nohighlight">
 \[P(arm) = \frac{ e ^  \frac{\mu_i - \max{\mu}}{ \tau } }
 { \Sigma{e ^  \frac{\mu - \max{\mu}}{ \tau }}  }\]</div>
-<p>where <span class="math notranslate nohighlight">\(\mu_i\)</span> is the mean reward for that arm and <span class="math notranslate nohighlight">\(\tau\)</span> is the “temperature” to determine the degree of
-exploration.</p>
+<p>where <span class="math notranslate nohighlight">\(\mu_i\)</span> is the mean reward for that arm and <span class="math notranslate nohighlight">\(\tau\)</span> is the “temperature” to determine
+the degree of exploration.</p>
 <dl class="py attribute">
 <dt class="sig sig-object py" id="mabwiser.mab.LearningPolicy.Softmax.tau">
 <span class="sig-name descname"><span class="pre">tau</span></span><a class="headerlink" href="#mabwiser.mab.LearningPolicy.Softmax.tau" title="Permalink to this definition"></a></dt>
@@ -885,7 +879,7 @@
 </div>
 <dl class="py method">
 <dt class="sig sig-object py" id="mabwiser.mab.MAB.add_arm">
-<span class="sig-name descname"><span class="pre">add_arm</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">arm</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Arm</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">binarizer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">scaler</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">NoReturn</span></span></span><a class="headerlink" href="#mabwiser.mab.MAB.add_arm" title="Permalink to this definition"></a></dt>
+<span class="sig-name descname"><span class="pre">add_arm</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">arm</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Arm</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">binarizer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">NoReturn</span></span></span><a class="headerlink" href="#mabwiser.mab.MAB.add_arm" title="Permalink to this definition"></a></dt>
 <dd><p>Adds an _arm_ to the list of arms.</p>
 <p>Incorporates the arm into the learning and neighborhood policies with no training data.</p>
 <dl class="field-list simple">
@@ -893,7 +887,6 @@
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>arm</strong> (<em>Arm</em>) – The new arm to be added.</p></li>
 <li><p><strong>binarizer</strong> (<em>Callable</em>) – The new binarizer function for Thompson Sampling.</p></li>
-<li><p><strong>scaler</strong> (<em>Callable</em>) – A scaler object from sklearn.preprocessing.</p></li>
 </ul>
 </dd>
 <dt class="field-even">Return type</dt>
@@ -902,8 +895,6 @@
 <dt class="field-odd">Raises</dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>TypeError</strong> – For ThompsonSampling, binarizer must be a callable function.:</p></li>
-<li><p><strong>TypeError</strong> – The standard scaler object must have a transform method.:</p></li>
-<li><p><strong>TypeError</strong> – The standard scaler object must be fit with calculated <code class="docutils literal notranslate"><span class="pre">mean_</span></code> and <code class="docutils literal notranslate"><span class="pre">var_</span></code> attributes.:</p></li>
 <li><p><strong>ValueError</strong> – A binarizer function was provided but the learning policy is not Thompson Sampling.:</p></li>
 <li><p><strong>ValueError</strong> – The arm already exists.:</p></li>
 <li><p><strong>ValueError</strong> – The arm is <code class="docutils literal notranslate"><span class="pre">None</span></code>.:</p></li>
@@ -1022,8 +1013,9 @@
 In testing, they return the best arm given new context(s).</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>contexts</strong> (<em>Union</em><em>[</em><em>None</em><em>, </em><em>List</em><em>[</em><em>List</em><em>[</em><em>Num</em><em>]</em><em>]</em><em>, </em><em>np.ndarray</em><em>, </em><em>pd.Series</em><em>, </em><em>pd.DataFrame</em><em>]</em>) – The context under which each decision is made. Default value is None.
-Contexts should be <code class="docutils literal notranslate"><span class="pre">None</span></code> for context-free bandits and is required for contextual bandits.</p>
+<dd class="field-odd"><p><strong>contexts</strong> (<em>Union</em><em>[</em><em>None</em><em>, </em><em>List</em><em>[</em><em>Num</em><em>]</em><em>, </em><em>List</em><em>[</em><em>List</em><em>[</em><em>Num</em><em>]</em><em>]</em><em>, </em><em>np.ndarray</em><em>, </em><em>pd.Series</em><em>, </em><em>pd.DataFrame</em><em>]</em>) – The context for the expected rewards. Default value is None.
+If contexts is not <code class="docutils literal notranslate"><span class="pre">None</span></code> for context-free bandits, the predictions returned will be a
+list of the same length as contexts.</p>
 </dd>
 <dt class="field-even">Return type</dt>
 <dd class="field-even"><p>The recommended arm or recommended arms list.</p>
@@ -1031,8 +1023,7 @@
 <dt class="field-odd">Raises</dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>TypeError</strong> – Contexts is not given as <code class="docutils literal notranslate"><span class="pre">None</span></code>, list, numpy array, pandas series or data frames.:</p></li>
-<li><p><strong>ValueError</strong> – Predicting with contexts data when there is no contextual policy.:</p></li>
-<li><p><strong>ValueError</strong> – Contextual policy when predicting with no contexts data.:</p></li>
+<li><p><strong>ValueError</strong> – Prediction with context policy requires context data.:</p></li>
 </ul>
 </dd>
 </dl>
@@ -1046,7 +1037,8 @@
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><p><strong>contexts</strong> (<em>Union</em><em>[</em><em>None</em><em>, </em><em>List</em><em>[</em><em>Num</em><em>]</em><em>, </em><em>List</em><em>[</em><em>List</em><em>[</em><em>Num</em><em>]</em><em>]</em><em>, </em><em>np.ndarray</em><em>, </em><em>pd.Series</em><em>, </em><em>pd.DataFrame</em><em>]</em>) – The context for the expected rewards. Default value is None.
-Contexts should be <code class="docutils literal notranslate"><span class="pre">None</span></code> for context-free bandits and is required for contextual bandits.</p>
+If contexts is not <code class="docutils literal notranslate"><span class="pre">None</span></code> for context-free bandits, the predicted expectations returned will be a
+list of the same length as contexts.</p>
 </dd>
 <dt class="field-even">Return type</dt>
 <dd class="field-even"><p>The dictionary of arms (key) to their expected rewards (value), or a list of such dictionaries.</p>
@@ -1054,8 +1046,7 @@
 <dt class="field-odd">Raises</dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>TypeError</strong> – Contexts is not given as <code class="docutils literal notranslate"><span class="pre">None</span></code>, list, numpy array or pandas data frames.:</p></li>
-<li><p><strong>ValueError</strong> – Predicting with contexts data when there is no contextual policy.:</p></li>
-<li><p><strong>ValueError</strong> – Contextual policy when predicting with no contexts data.:</p></li>
+<li><p><strong>ValueError</strong> – Prediction with context policy requires context data.:</p></li>
 </ul>
 </dd>
 </dl>
@@ -1286,7 +1277,7 @@
 <dl class="py attribute">
 <dt class="sig sig-object py" id="mabwiser.mab.NeighborhoodPolicy.LSHNearest.no_nhood_prob_of_arm">
 <span class="sig-name descname"><span class="pre">no_nhood_prob_of_arm</span></span><a class="headerlink" href="#mabwiser.mab.NeighborhoodPolicy.LSHNearest.no_nhood_prob_of_arm" title="Permalink to this definition"></a></dt>
-<dd><p>The probabilities associated with each arm. Used to select random arm if a prediction context has no neighbors.
+<dd><p>The probabilities associated with each arm. Used to select random arm if context has no neighbors.
 If not given, a uniform random distribution over all arms is assumed.
 The probabilities should sum up to 1.</p>
 <dl class="field-list simple">
@@ -1364,7 +1355,7 @@
 <dl class="py attribute">
 <dt class="sig sig-object py" id="mabwiser.mab.NeighborhoodPolicy.Radius.no_nhood_prob_of_arm">
 <span class="sig-name descname"><span class="pre">no_nhood_prob_of_arm</span></span><a class="headerlink" href="#mabwiser.mab.NeighborhoodPolicy.Radius.no_nhood_prob_of_arm" title="Permalink to this definition"></a></dt>
-<dd><p>The probabilities associated with each arm. Used to select random arm if a prediction context has no neighbors.
+<dd><p>The probabilities associated with each arm. Used to select random arm if context has no neighbors.
 If not given, a uniform random distribution over all arms is assumed.
 The probabilities should sum up to 1.</p>
 <dl class="field-list simple">
diff --git a/docs/contributing.html b/docs/contributing.html
index 1c350f2..51d6959 100644
--- a/docs/contributing.html
+++ b/docs/contributing.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Contributing &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Contributing &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/examples.html b/docs/examples.html
index 2d9b444..5c46ef8 100644
--- a/docs/examples.html
+++ b/docs/examples.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Usage Examples &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Usage Examples &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/genindex.html b/docs/genindex.html
index acc6690..ade722a 100644
--- a/docs/genindex.html
+++ b/docs/genindex.html
@@ -3,7 +3,7 @@
 <head>
   <meta charset="utf-8" />
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Index &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Index &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
@@ -111,21 +111,13 @@ <h2 id="A">A</h2>
       <li><a href="api.html#mabwiser.utils.argmax">argmax() (in module mabwiser.utils)</a>
 </li>
       <li><a href="api.html#mabwiser.utils.argmin">argmin() (in module mabwiser.utils)</a>
-</li>
-      <li><a href="api.html#mabwiser.utils.Arm">Arm (in module mabwiser.utils)</a>
-</li>
-      <li><a href="api.html#mabwiser.base_mab.BaseMAB.arm_to_expectation">arm_to_expectation (mabwiser.base_mab.BaseMAB attribute)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
-      <li><a href="api.html#id1">arm_to_scaler (mabwiser.mab.LearningPolicy.LinGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinGreedy.arm_to_scaler">[1]</a>
-
-      <ul>
-        <li><a href="api.html#id5">(mabwiser.mab.LearningPolicy.LinTS attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinTS.arm_to_scaler">[1]</a>
+      <li><a href="api.html#mabwiser.utils.Arm">Arm (in module mabwiser.utils)</a>
 </li>
-        <li><a href="api.html#id8">(mabwiser.mab.LearningPolicy.LinUCB attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinUCB.arm_to_scaler">[1]</a>
+      <li><a href="api.html#mabwiser.base_mab.BaseMAB.arm_to_expectation">arm_to_expectation (mabwiser.base_mab.BaseMAB attribute)</a>
 </li>
-      </ul></li>
       <li><a href="api.html#mabwiser.simulator.Simulator.arm_to_stats_test">arm_to_stats_test (mabwiser.simulator.Simulator attribute)</a>
 </li>
       <li><a href="api.html#mabwiser.simulator.Simulator.arm_to_stats_total">arm_to_stats_total (mabwiser.simulator.Simulator attribute)</a>
@@ -223,7 +215,7 @@ <h2 id="E">E</h2>
       <li><a href="api.html#id0">epsilon (mabwiser.mab.LearningPolicy.EpsilonGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.EpsilonGreedy.epsilon">[1]</a>
 
       <ul>
-        <li><a href="api.html#id2">(mabwiser.mab.LearningPolicy.LinGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinGreedy.epsilon">[1]</a>
+        <li><a href="api.html#id1">(mabwiser.mab.LearningPolicy.LinGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinGreedy.epsilon">[1]</a>
 </li>
       </ul></li>
   </ul></td>
@@ -284,12 +276,12 @@ <h2 id="K">K</h2>
 <h2 id="L">L</h2>
 <table style="width: 100%" class="indextable genindextable"><tr>
   <td style="width: 33%; vertical-align: top;"><ul>
-      <li><a href="api.html#id3">l2_lambda (mabwiser.mab.LearningPolicy.LinGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinGreedy.l2_lambda">[1]</a>
+      <li><a href="api.html#id2">l2_lambda (mabwiser.mab.LearningPolicy.LinGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinGreedy.l2_lambda">[1]</a>
 
       <ul>
-        <li><a href="api.html#id6">(mabwiser.mab.LearningPolicy.LinTS attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinTS.l2_lambda">[1]</a>
+        <li><a href="api.html#id5">(mabwiser.mab.LearningPolicy.LinTS attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinTS.l2_lambda">[1]</a>
 </li>
-        <li><a href="api.html#id9">(mabwiser.mab.LearningPolicy.LinUCB attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinUCB.l2_lambda">[1]</a>
+        <li><a href="api.html#id8">(mabwiser.mab.LearningPolicy.LinUCB attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinUCB.l2_lambda">[1]</a>
 </li>
       </ul></li>
       <li><a href="api.html#mabwiser.mab.MAB.learning_policy">learning_policy (mabwiser.mab.MAB attribute)</a>
@@ -483,10 +475,18 @@ <h2 id="R">R</h2>
 <h2 id="S">S</h2>
 <table style="width: 100%" class="indextable genindextable"><tr>
   <td style="width: 33%; vertical-align: top;"><ul>
-      <li><a href="api.html#mabwiser.simulator.Simulator.scaler">scaler (mabwiser.simulator.Simulator attribute)</a>
+      <li><a href="api.html#id3">scale (mabwiser.mab.LearningPolicy.LinGreedy attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinGreedy.scale">[1]</a>
+
+      <ul>
+        <li><a href="api.html#id6">(mabwiser.mab.LearningPolicy.LinTS attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinTS.scale">[1]</a>
 </li>
+        <li><a href="api.html#id9">(mabwiser.mab.LearningPolicy.LinUCB attribute)</a>, <a href="api.html#mabwiser.mab.LearningPolicy.LinUCB.scale">[1]</a>
+</li>
+      </ul></li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
+      <li><a href="api.html#mabwiser.simulator.Simulator.scaler">scaler (mabwiser.simulator.Simulator attribute)</a>
+</li>
       <li><a href="api.html#mabwiser.mab.MAB.seed">seed (mabwiser.mab.MAB attribute)</a>
 </li>
       <li><a href="api.html#mabwiser.simulator.Simulator">Simulator (class in mabwiser.simulator)</a>
diff --git a/docs/index.html b/docs/index.html
index d074e6f..1fa6174 100644
--- a/docs/index.html
+++ b/docs/index.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>MABWiser Contextual Multi-Armed Bandits &mdash; MABWiser 2.3.0 documentation</title>
+  <title>MABWiser Contextual Multi-Armed Bandits &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/installation.html b/docs/installation.html
index 7a3917d..4950194 100644
--- a/docs/installation.html
+++ b/docs/installation.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Installation &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Installation &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/new_bandit.html b/docs/new_bandit.html
index 9cdd760..5c14660 100644
--- a/docs/new_bandit.html
+++ b/docs/new_bandit.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Adding a New Bandit &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Adding a New Bandit &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/objects.inv b/docs/objects.inv
index 9d6a0d4..706ac53 100644
Binary files a/docs/objects.inv and b/docs/objects.inv differ
diff --git a/docs/py-modindex.html b/docs/py-modindex.html
index 5c042c8..9286a02 100644
--- a/docs/py-modindex.html
+++ b/docs/py-modindex.html
@@ -3,7 +3,7 @@
 <head>
   <meta charset="utf-8" />
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Python Module Index &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Python Module Index &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/quick.html b/docs/quick.html
index d2b7de1..3a8e159 100644
--- a/docs/quick.html
+++ b/docs/quick.html
@@ -4,7 +4,7 @@
   <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Quick Start &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Quick Start &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <!--[if lt IE 9]>
diff --git a/docs/search.html b/docs/search.html
index f4e4735..792b41e 100644
--- a/docs/search.html
+++ b/docs/search.html
@@ -3,7 +3,7 @@
 <head>
   <meta charset="utf-8" />
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Search &mdash; MABWiser 2.3.0 documentation</title>
+  <title>Search &mdash; MABWiser 2.4.0 documentation</title>
       <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
       <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
     
diff --git a/docs/searchindex.js b/docs/searchindex.js
index da13c11..62584d2 100644
--- a/docs/searchindex.js
+++ b/docs/searchindex.js
@@ -1 +1 @@
-Search.setIndex({docnames:["about","api","contributing","examples","index","installation","new_bandit","quick"],envversion:{"sphinx.domains.c":2,"sphinx.domains.changeset":1,"sphinx.domains.citation":1,"sphinx.domains.cpp":4,"sphinx.domains.index":1,"sphinx.domains.javascript":2,"sphinx.domains.math":2,"sphinx.domains.python":3,"sphinx.domains.rst":2,"sphinx.domains.std":2,sphinx:56},filenames:["about.rst","api.rst","contributing.rst","examples.rst","index.rst","installation.rst","new_bandit.rst","quick.rst"],objects:{"mabwiser.base_mab":[[1,1,1,"","BaseMAB"]],"mabwiser.base_mab.BaseMAB":[[1,2,1,"","add_arm"],[1,3,1,"","arm_to_expectation"],[1,3,1,"","arms"],[1,3,1,"","backend"],[1,3,1,"","cold_arm_to_warm_arm"],[1,2,1,"","fit"],[1,3,1,"","n_jobs"],[1,2,1,"","partial_fit"],[1,2,1,"","predict"],[1,2,1,"","predict_expectations"],[1,2,1,"","remove_arm"],[1,3,1,"","rng"],[1,3,1,"","trained_arms"],[1,2,1,"","warm_start"]],"mabwiser.mab":[[1,1,1,"","LearningPolicy"],[1,1,1,"","MAB"],[1,1,1,"","NeighborhoodPolicy"]],"mabwiser.mab.LearningPolicy":[[1,1,1,"","EpsilonGreedy"],[1,1,1,"","LinGreedy"],[1,1,1,"","LinTS"],[1,1,1,"","LinUCB"],[1,1,1,"","Popularity"],[1,1,1,"","Random"],[1,1,1,"","Softmax"],[1,1,1,"","ThompsonSampling"],[1,1,1,"","UCB1"]],"mabwiser.mab.LearningPolicy.EpsilonGreedy":[[1,3,1,"id0","epsilon"]],"mabwiser.mab.LearningPolicy.LinGreedy":[[1,3,1,"id1","arm_to_scaler"],[1,3,1,"id2","epsilon"],[1,3,1,"id3","l2_lambda"]],"mabwiser.mab.LearningPolicy.LinTS":[[1,3,1,"id4","alpha"],[1,3,1,"id5","arm_to_scaler"],[1,3,1,"id6","l2_lambda"]],"mabwiser.mab.LearningPolicy.LinUCB":[[1,3,1,"id7","alpha"],[1,3,1,"id8","arm_to_scaler"],[1,3,1,"id9","l2_lambda"]],"mabwiser.mab.LearningPolicy.Softmax":[[1,3,1,"id10","tau"]],"mabwiser.mab.LearningPolicy.ThompsonSampling":[[1,3,1,"id11","binarizer"]],"mabwiser.mab.LearningPolicy.UCB1":[[1,3,1,"id12","alpha"]],"mabwiser.mab.MAB":[[1,2,1,"","add_arm"],[1,3,1,"","arms"],[1,3,1,"","backend"],[1,2,1,"","fit"],[1,3,1,"","is_contextual"],[1,4,1,"id13","learning_policy"],[1,3,1,"","n_jobs"],[1,4,1,"id14","neighborhood_policy"],[1,2,1,"","partial_fit"],[1,2,1,"","predict"],[1,2,1,"","predict_expectations"],[1,2,1,"","remove_arm"],[1,3,1,"","seed"],[1,2,1,"","warm_start"]],"mabwiser.mab.NeighborhoodPolicy":[[1,1,1,"","Clusters"],[1,1,1,"","KNearest"],[1,1,1,"","LSHNearest"],[1,1,1,"","Radius"],[1,1,1,"","TreeBandit"]],"mabwiser.mab.NeighborhoodPolicy.Clusters":[[1,3,1,"id15","is_minibatch"],[1,3,1,"id16","n_clusters"]],"mabwiser.mab.NeighborhoodPolicy.KNearest":[[1,3,1,"id17","k"],[1,3,1,"id18","metric"]],"mabwiser.mab.NeighborhoodPolicy.LSHNearest":[[1,3,1,"id19","n_dimensions"],[1,3,1,"id20","n_tables"],[1,3,1,"id21","no_nhood_prob_of_arm"]],"mabwiser.mab.NeighborhoodPolicy.Radius":[[1,3,1,"id22","metric"],[1,3,1,"id23","no_nhood_prob_of_arm"],[1,3,1,"id24","radius"]],"mabwiser.mab.NeighborhoodPolicy.TreeBandit":[[1,3,1,"id27","tree_parameters"]],"mabwiser.simulator":[[1,1,1,"","Simulator"],[1,5,1,"","default_evaluator"]],"mabwiser.simulator.Simulator":[[1,3,1,"","arm_to_stats_test"],[1,3,1,"","arm_to_stats_total"],[1,3,1,"","arm_to_stats_train"],[1,3,1,"","arms"],[1,3,1,"","bandit_to_arm_to_stats_avg"],[1,3,1,"","bandit_to_arm_to_stats_max"],[1,3,1,"","bandit_to_arm_to_stats_min"],[1,3,1,"","bandit_to_arm_to_stats_neighborhoods"],[1,3,1,"","bandit_to_confusion_matrices"],[1,3,1,"","bandit_to_expectations"],[1,3,1,"","bandit_to_neighborhood_size"],[1,3,1,"","bandit_to_predictions"],[1,3,1,"","bandits"],[1,3,1,"","batch_size"],[1,3,1,"","contexts"],[1,3,1,"","decisions"],[1,3,1,"","evaluator"],[1,2,1,"","get_arm_stats"],[1,2,1,"","get_stats"],[1,3,1,"","is_ordered"],[1,3,1,"","is_quick"],[1,3,1,"","logger"],[1,2,1,"","plot"],[1,3,1,"","rewards"],[1,2,1,"","run"],[1,3,1,"","scaler"],[1,3,1,"","test_indices"],[1,3,1,"","test_size"]],"mabwiser.utils":[[1,6,1,"","Arm"],[1,1,1,"","Constants"],[1,6,1,"","Num"],[1,5,1,"","argmax"],[1,5,1,"","argmin"],[1,5,1,"","check_false"],[1,5,1,"","check_true"],[1,5,1,"","create_rng"],[1,5,1,"","reset"]],"mabwiser.utils.Constants":[[1,3,1,"","default_seed"],[1,3,1,"","distance_metrics"]],mabwiser:[[1,0,0,"-","base_mab"],[1,0,0,"-","mab"],[1,0,0,"-","simulator"],[1,0,0,"-","utils"]]},objnames:{"0":["py","module","Python module"],"1":["py","class","Python class"],"2":["py","method","Python method"],"3":["py","attribute","Python attribute"],"4":["py","property","Python property"],"5":["py","function","Python function"],"6":["py","data","Python data"]},objtypes:{"0":"py:module","1":"py:class","2":"py:method","3":"py:attribute","4":"py:property","5":"py:function","6":"py:data"},terms:{"0":[1,3,6],"00129":4,"01":3,"05":1,"1":[0,1,3,4,5,7],"10":[1,3,4],"100":3,"1000":3,"100000":3,"11":[1,3],"1109":4,"111":3,"1142":4,"12":3,"123456":[1,3,6],"13":3,"15":[1,3],"17":[1,3,4,6,7],"18":3,"19":[3,4],"1d":3,"2":[0,1,3,4,5],"20":[1,3,4,6,7],"200":3,"2010":0,"2016":0,"2017":1,"2019":4,"2021":4,"21":3,"2150021":4,"22":3,"23":3,"25":[1,3,4,6,7],"26":3,"27":3,"2d":3,"3":[1,3,4,5],"30":[1,4],"31st":4,"33":3,"37":3,"38":3,"39":3,"4":[1,3,4],"42":[3,6],"48":3,"5":[1,3],"50":3,"52":3,"55":3,"56":3,"57":3,"6":[3,4,5],"65":3,"68":3,"7":[1,3],"75":3,"8":[2,3,4],"83":3,"8995418":[],"9":[1,3,4,6,7],"909":4,"914":4,"99":3,"abstract":[1,6],"boolean":1,"break":6,"case":[1,3,6],"class":[1,6],"default":[1,2,6],"do":[0,1,2,6],"final":6,"float":[1,6],"function":[1,2,3,6],"import":[0,1,3,4,5,6,7],"int":[1,4,6],"long":0,"new":[0,1,2,3,4],"public":[2,4],"return":[1,6],"static":1,"super":6,"true":[1,3,6],"try":5,A:[0,1,3,4,6],As:[0,1],At:0,For:[0,1,5,6],If:[1,5,6],In:[0,1,2,3,6],It:[1,3,4,6],No:1,Not:1,OR:4,Or:6,That:6,The:[0,1,2,3,4,5,6],Then:1,There:[0,3,5],These:[1,6],To:[1,5,6],_:[3,6],__init__:[1,6],__version__:6,_arm_:1,_baserng:[1,6],_copy_arm:6,_fit_arm:6,_mycoolpolici:6,_onlin:1,_parallel_fit:6,_predict_context:6,_uptake_new_arm:[1,6],_valid:6,_validate_mab_arg:6,ab:[],abl:[0,6],about:[1,3,4,6],abov:6,accept:1,access:[1,6],accident:6,accompani:2,accord:6,accordingli:1,action:[0,6],activ:6,actual:1,ad:[1,3,4],adam:1,add:[1,6],add_arm:[1,3,6],addit:[1,3,6],additional_layout:3,additional_revenu:3,adher:[2,4],adjust:1,advantag:6,advertis:[0,3],affect:2,after:[0,1],ag:3,agent:0,aggreg:1,al:0,algorithm:[0,1,4],alia:1,all:[1,2,5,6],allow:[3,6],alpha:[1,3,4,7],alreadi:1,also:[0,1],altern:5,among:[0,3],amount:3,an:[0,1,3,4,6,7],ani:[1,2,5,6],api:[2,4],append:3,appli:[1,6],applic:[0,2],approach:1,approxim:1,ar:[0,1,2,3,5,6],arang:3,argmax:1,argmin:1,argument:1,arm1:[1,4,6,7],arm2:[1,4,6,7],arm3:1,arm:[1,3,6,7],arm_to_expect:[1,6],arm_to_featur:[1,3,6],arm_to_predict:6,arm_to_scal:1,arm_to_stat:1,arm_to_stats_test:1,arm_to_stats_tot:1,arm_to_stats_train:1,arm_to_threshold:1,arms_to_stat:1,arms_to_stats_train:1,arrai:[1,3,6],articl:4,artif:4,artifici:4,arxiv:0,asctim:1,assert:[3,6],assertdictequ:6,assertequ:6,assertionerror:1,assess:6,assign:1,associ:1,assum:1,assumpt:1,attribut:[1,6],author:4,automat:6,avail:[1,3,6],averag:[1,3],avg:[1,3],b:[3,4,6],back:6,backend:[1,6],bandit:[1,3],bandit_to_arm_to_stat:1,bandit_to_arm_to_stats_avg:[1,3],bandit_to_arm_to_stats_max:[1,3],bandit_to_arm_to_stats_min:[1,3],bandit_to_arm_to_stats_neighborhood:1,bandit_to_confusion_matric:1,bandit_to_expect:1,bandit_to_neighborhood_s:1,bandit_to_predict:1,base:[1,3,4,6,7],base_mab:4,basemab:[1,6],basetest:6,basic:[1,2],batch:[1,3],batch_siz:[1,3],bdist_wheel:5,becom:3,been:[1,6],befor:[1,6],behav:6,behavior:[1,6],behind:[3,6],being:1,below:3,bernard:4,best:[0,1,3,6],best_arm:6,beta:1,between:[0,1,4,6,7],bibtex:[],binar:[1,6],binari:1,block:[],booktitl:4,bool:[1,6],both:4,bound1:1,bound:[1,3,4],branch:6,braycurti:1,build:5,built:[3,4],cach:5,calcul:[1,6],call:[0,1,6],callabl:[1,6],can:[0,1,3,4,5,6],canberra:1,cannot:[1,6],capabl:3,captur:1,carnegi:0,cd:5,cdist:1,center:4,certain:1,chang:6,changelog:6,chebyshev:1,check:[1,2,6],check_fals:1,check_tru:1,checkout:6,choic:3,choleski:1,choos:[0,1,7],chosen:1,chronolog:1,cite:4,cityblock:1,clean:2,click:3,click_rat:3,clinic:0,clone:5,closest:1,cluster:[1,4],code:[1,2,6],coeffici:1,coher:2,cold:[1,6],cold_arm:6,cold_arm_to_warm_arm:[1,6],collect:1,collis:1,column:1,com:5,commerc:3,common:1,commun:1,compar:1,comparison:4,compat:1,complet:[1,6],compon:4,computation:1,conceptu:3,concurr:1,conf:4,confer:4,confid:[1,3,4],confirm:5,confus:1,congratul:6,connect:6,conserv:1,consid:6,consist:6,constant:1,constructor:[1,6],contain:[1,2],context:[0,1,4,6],contexts_test:3,contexts_train:3,contextu:[0,1],contextual_mab:1,continu:6,contribut:[4,6],control:1,convert:1,cool:6,copi:6,core:[1,3,6],correct:1,correctli:6,correl:1,correspond:[1,3,6],cosin:1,count:1,counter:6,cours:0,covari:1,cover:6,cpu:1,creat:[1,6],create_rng:1,cumul:[0,1],custom:[1,3,5],d:1,danc:6,data:[1,3,4,6,7],datafram:[1,3],dataset:3,dblp:4,decid:[0,6],decis:[0,1,3,4,6,7],decisions_test:3,decisions_train:3,decisiontreeclassifi:1,declar:[1,6],decomposit:1,decor:6,decreas:1,deem:1,deepcopi:6,def:[1,6],default_evalu:1,default_se:1,defin:[1,6],definit:[0,1],degre:1,denot:[1,6],depend:[1,3],descript:[1,6],deserv:6,design:[2,3],detail:6,determin:[1,6],determinist:[0,1],develop:6,dice:1,dict:[1,6],dictat:6,dictionari:[1,6],differ:[0,1,3,6],dimens:1,dimension:1,dir:5,directli:[1,6],directori:2,discov:5,displai:3,dist:5,distanc:[1,6],distance_metr:1,distance_quantil:[1,3,6],distribut:[0,1],divis:1,doc:6,docsrc:6,doctest:6,document:[2,6],doe:[1,6],doi:4,done:6,down:6,dure:[3,6],dxd:1,e:[1,3,4,6],each:[0,1,3,6],easiest:2,effect:[1,6],eg:1,elmachtoub:1,emerg:6,emili:4,empti:1,enabl:6,engin:0,ensur:[1,2],entir:1,entri:[1,6],epsilon:[1,3,4],epsilongreedi:[1,3],equal:3,error:6,estim:1,et:0,etc:1,euclidean:1,evalu:1,even:6,event:3,everi:[0,1,3,6],everyon:6,evolv:6,exampl:[1,4,5,6,7],excel:[0,4],except:1,exchang:1,execut:[1,6],exist:[0,1,2,6],exp:6,expect:[0,1,3,4,6,7],expens:1,experi:[1,3],explor:[1,3],expos:4,express:[1,6],extend:[1,6],face:0,factor:1,fals:[1,3,6],fashion:6,fast:[1,4],featur:[1,2,3,4,6],fidel:[4,5],field:[1,3,6],file:6,filter:1,first:[1,6],fit:[1,3,4,6,7],fit_arm:6,fit_transform:3,flag:[1,6],floot:[],folder:[3,5,6],follow:[1,4,6],foremost:6,format:1,formula:6,frac:1,frame:[1,3],framework:6,free:[1,4],from:[0,1,3,4,5,6,7],fromkei:6,g:[1,6],gener:[1,2,3,6],get:[1,2],get_arm_stat:1,get_stat:1,git:[5,6],github:[2,4,5,6],give:1,given:[0,1,3,6],global:1,go:6,goe:[1,6],good:1,greater:1,greedi:[1,3,4,6],greedili:1,guarante:1,guidelin:6,ha:[0,1,6],ham:1,handl:6,happi:2,hash:1,have:[0,1,2,6],heavili:4,helper:1,here:[3,6],high:3,higher:1,highest:[0,1],hint:6,histor:[1,3],histori:[1,3],homepag:3,host:[2,4],how:[1,3,4,5,6,7],http:[4,5],hyper:[1,3,4,6],hyper_parameter_tun:3,hyperplan:1,i:[1,3],i_d:1,ictai:4,idea:[3,6],ident:1,ieee:4,ieeexplor:[],ijait:4,imagin:6,immut:6,implement:[1,4],implementor:[1,6],includ:[5,6],incorpor:1,increas:1,increment:6,independ:3,index:[1,4],indic:1,induc:1,infin:1,inform:[0,3],inherit:[1,6],initi:[1,6],inner:6,inproceed:4,input:[1,6],instal:[4,6],instead:[1,6],integ:1,intel:4,intellig:4,intend:1,interfac:[1,6],intern:[1,4,6],internet:0,interpret:1,introduc:[1,6],invers:1,invest:4,involv:6,is_contextu:1,is_minibatch:1,is_ord:[1,3],is_per_arm:[1,3],is_predict:6,is_quick:1,isn:6,issu:[2,4],item:[1,6],its:[1,6],j:4,jaccard:1,joblib:1,journal:4,k:[0,1,4],kadioglu:4,keep:1,kei:1,keyword:6,kleynhan:4,kmean:1,knearest:1,know:0,kulsinski:1,kwarg:1,l2_lambda:[1,3],lambda:1,larg:3,larger:[],later:6,latest:5,layout:3,lead:3,leaf:1,learn:[0,1,3,4,6,7],learning_polici:[1,3,6],learningpolici:[1,3,4,5,6,7],learnt:3,least:1,leav:1,length:1,less:1,let:6,level:3,levelnam:1,leverag:6,li:0,librari:[1,4,6,7],like:6,likelihood:1,linear:3,lingreedi:[1,4],lint:[1,4],linucb:[1,3,4],list:[1,3,5,6],list_of_arm:1,listen:3,live:6,local:[1,6],lock:1,log:1,log_fil:1,log_format:1,logger:1,logist:1,loki:1,look:6,lot:1,low:1,lp:6,lsh:4,lshnearest:1,lu:0,m:5,mab1:1,mab2:1,mab:[0,4,5,6,7],mab_nam:3,mabwis:[2,3,5,6,7],machin:0,made:1,magic:6,mahalanobi:1,mai:0,main:[2,6],make:[0,1,2,3,4,6],make_classif:3,mani:[0,1],map:[1,6],marek:1,master:[5,6],match:1,math:6,matric:1,matrix:1,max:1,maxim:0,maximum:1,mcnelli:1,md:6,mean:[1,3],mean_:1,mellon:0,member:6,memori:1,merg:6,messag:1,meta:6,method:[1,2,3,6],metric:[1,3],might:6,min:1,minibatchkmean:1,minimum:[1,6],minkowski:1,mismatch:1,miss:1,mode:6,model:[0,1,3,4,6,7],model_select:3,model_to_confusion_matric:1,model_to_predict:1,models_to_evalu:1,modif:2,modifi:1,modul:[1,4],more:[0,1],most:6,move:6,mu:1,mu_i:1,much:0,multi:1,multipl:[0,1,3],multipli:1,multiprocess:1,multivari:1,music:3,must:1,my_paramet:6,my_results_:3,my_value_to_arm:6,mycoolbandittest:6,mycoolpolici:6,n:[1,3],n_cluster:1,n_context_col:1,n_dimens:1,n_featur:3,n_i:1,n_inform:3,n_job:[1,3,6],n_row:1,n_sampl:3,n_tabl:1,name:[1,6],namedtupl:6,nan:1,nd:1,ndarrai:[1,6],nearest:[1,4],necessari:[5,6],need:[0,3,6],neg:[1,6],neighbor:1,neighborhood:[1,3,4,6],neighborhood_polici:[1,3],neighborhoodpolici:[1,3,4,5,6,7],next:[3,6],nn:1,no_nhood_prob_of_arm:1,non:[1,4],none:[1,3,5,6],noreturn:[1,6],normal:1,note:[1,6],notebook:5,notic:[1,6],notimplementederror:1,novemb:4,now:6,np:[1,3,6],nparrai:1,num:[1,6],num_run:6,number:[1,4,6],numer:1,numpi:[1,3,6],numpydoc:2,o:1,object:[0,1,6],observ:[0,1,3,6],obtain:1,offlin:[1,3],offline_sim:1,oh:1,older:2,one:[1,6],onli:[0,1],onlin:[0,1,3,6],onto:[1,6],oper:[1,6],optim:3,option:[0,1,3,6],order:[1,3],org:4,origin:5,other:[0,1,3,6],otherwis:1,our:[2,6],out:[1,2],outcom:0,output:1,over:1,overal:0,overhead:1,overrid:1,own:6,p:1,packag:[5,6],page:4,pairwis:6,panda:[1,3],parallel:[1,4,6],paralleliz:4,param:6,paramet:[1,3,4,6],parametr:[1,4],part:[0,6],partial:1,partial_fit:[1,3,6],partit:1,pass:[2,5,6],past:3,pd:[1,3],pep:[2,4],per:3,per_arm:1,perform:[1,6],petrik:1,phase:[3,6],pip:[5,6],pipelin:1,plai:0,platform:5,playlist:3,pleas:[4,6],plot:[1,3],point:[1,6],polici:[1,3,6,7],pool:1,popular:[1,4,6],portland:4,posit:1,possibl:[3,6],potenti:[0,1],practic:1,precis:1,predict:[1,3,4,6,7],predict_expect:[1,3,6],prefix:[3,6],prepar:6,preprocess:[1,3],present:1,previou:[1,2,3,6],print:3,privat:6,probabl:1,probe:3,problem:[0,1,3],proc:0,procedur:1,process:1,progress:1,project:[1,2,5],promot:1,properli:6,properti:[1,6],proportion:1,prototyp:[1,4],provid:[1,4,6],publish:4,pull:[0,4,5],py3:5,py:[5,6],pydoc:6,python:[1,4,5,6],quantil:1,question:0,quickli:6,radiu:[1,3,4],rais:[1,6],randint:3,random:[1,3,4,6],random_st:3,randomli:1,rang:3,rate:3,rather:1,ration:1,re:3,read:1,readi:6,readm:6,real:[0,6],realiti:3,receiv:[0,1],recommend:[0,1,3],recompil:6,refer:[0,5],reflect:1,region:1,regist:6,regress:1,regular:1,releas:4,reli:1,remain:6,remov:[1,6],remove_arm:1,renew:0,repeat:1,repo:[2,3,5],represent:1,request:4,requir:[1,2,4],research:[0,1,4],reset:[1,6],resourc:0,result:[0,1,3],retriev:[1,3],revenu:3,review:6,reward:[0,1,3,4,6,7],rewards_test:3,rewards_train:3,ridg:1,rng:[1,6],robust:1,rogerstanimoto:1,root:1,routin:1,row:1,rst:6,run:[0,1,3,5,6],russellrao:1,ryan:1,s0218213021500214:4,s:[1,3,4,6],same:[1,6],sampl:[0,1,2,4],save:3,save_result:3,scale:3,scaler:[1,3,6],scenario:3,scikit:4,scipi:1,scratch:[5,6],sdist:5,search:6,sechan:1,section:[2,6],see:[1,6],seed:[1,3,6],seen:6,select:1,self:6,send:4,sensit:1,separ:1,sequenc:0,sequenti:0,serdar:4,seri:[1,3,6],servic:3,set:[0,1,3,6],setup:4,setuptool:5,seuclidean:1,share:[1,6],shell:5,should:[1,2,3,5],show:[3,4,7],shown:[3,6],side:0,sigma:1,sign:1,signatur:[1,6],sim:3,simhash:1,similar:6,simpl:[1,6],simul:4,sinc:3,singl:1,singular:[],situat:0,size:[1,3],skeleton:1,skip:[1,3],sklearn:[1,3],small:6,so:[1,6],softmax:[1,4],sokalmichen:1,sokalsneath:1,solv:[1,3],some:[0,1,6],someth:0,space:1,spatial:1,specif:1,specifi:1,speed:1,sphinx:6,split:[1,3],sqeuclidean:1,sqrt:1,squar:1,standard:[1,2,3,4],standardscal:3,start:[1,2,3,6],start_index:[1,6],stat:[1,3],statis:3,statist:1,std:1,step:[1,3,6],stick:0,still:2,stochast:0,store:[1,6],str:[1,3,6],stream:3,strength:1,strengthen:6,string:[1,6],strong:4,strongkk19:4,strongkk21:4,style:4,sub:1,subscrib:3,success:[1,5,6],suffer:1,suit:6,suitabl:0,sum:[1,6],support:[1,3,4],sure:[2,6],survei:0,svd:[],system:6,t:[0,1,6],tabl:1,take:[1,6],tau:1,tell:6,temperatur:1,test:[0,1,2,3,4,7],test_:6,test_add_new_arm:6,test_bas:6,test_df:3,test_df_revenu:3,test_fit_twic:6,test_indic:1,test_input_typ:6,test_invalid:6,test_my_paramet:6,test_mycoolbandit:6,test_parallel:6,test_partial_fit:6,test_remove_existing_arm:6,test_simple_usecase_arm:6,test_simple_usecase_expect:6,test_siz:[1,3],test_unused_arm:6,test_warm_start:6,test_within_neighborhood_polici:6,test_zero_reward:6,than:1,thei:[0,1,6],them:1,thi:[0,1,3,6],thompson:[1,4],thompsonsampl:1,those:6,thread:1,three:6,time:[0,1,6],titl:4,todo:6,togeth:1,toi:3,too:6,tool:4,total:[1,3,6],track:[1,4],train:[1,3,4,6,7],train_df:3,train_test_split:[1,3],trained_arm:1,transform:[1,3],tree:1,tree_paramet:1,treebandit:[1,4],treeheurist:1,trial:[0,1],ts:4,tune:[1,3,4],tupl:1,two:[4,5,6,7],txt:[5,6],tyler:0,type:[1,6],typeerror:1,typic:6,uai:1,ucb1:[1,3,4,7],ucb:1,uncertainti:0,under:[0,1,6],underli:[0,6],uniform:1,uniformli:1,uninstal:6,union:1,uniqu:1,unit:[2,6],unittest:[5,6],univers:0,unknown:0,unlik:6,untrain:[1,6],unus:6,up:1,updat:[1,3,6],upgrad:4,upper:[1,3,4],url:4,us:[1,2,3,4,5,6,7],usa:4,usag:[4,5,6],user:[3,6],util:[4,6],valid:[1,6],valu:[1,6],valueerror:1,var_:1,variant:0,vector:1,veri:1,version:[1,5,6],via:6,victori:6,violat:1,volum:4,wa:[1,5],wai:2,want:[2,3],warm:[1,3,6],warm_arm:6,warm_start:[1,3,6],we:[0,2,3,6],websit:3,weight:[1,3],were:6,what:[0,1,3,6],wheel:5,when:[1,6],where:[0,1,6],whether:[1,3],which:[0,1,3,6],whl:5,within:[1,6],work:[1,2,6],worker:1,world:[0,6],worldscientif:[],worst:3,would:6,wrapper:6,www:[],x:[1,5],x_i:1,year:4,yet:0,you:[1,2,3,4,5,6],your:[4,6],zero:[1,6],zhou:0},titles:["About Multi-Armed Bandits","MABWiser Public API","Contributing","Usage Examples","MABWiser Contextual Multi-Armed Bandits","Installation","Adding a New Bandit","Quick Start"],titleterms:{"1":6,"2":6,"3":6,"4":6,"new":6,"public":[1,6],about:0,ad:6,algorithm:6,api:[1,6],arm:[0,4],avail:4,bandit:[0,4,6],base_mab:1,bug:4,citat:4,code:[4,5],content:4,context:3,contextu:[3,4],contribut:2,exampl:3,exploit:0,explor:0,expos:6,free:3,high:6,implement:6,indic:4,instal:5,level:6,librari:5,mab:[1,3],mabwis:[1,4],multi:[0,4],non:3,option:5,overview:6,parallel:3,parametr:3,polici:4,pull:6,quick:[4,7],report:4,request:6,requir:5,send:6,setup:5,simul:[1,3],sourc:[4,5],start:[4,7],tabl:4,test:[5,6],upgrad:5,usag:3,util:1,vs:0,your:5}})
\ No newline at end of file
+Search.setIndex({docnames:["about","api","contributing","examples","index","installation","new_bandit","quick"],envversion:{"sphinx.domains.c":2,"sphinx.domains.changeset":1,"sphinx.domains.citation":1,"sphinx.domains.cpp":4,"sphinx.domains.index":1,"sphinx.domains.javascript":2,"sphinx.domains.math":2,"sphinx.domains.python":3,"sphinx.domains.rst":2,"sphinx.domains.std":2,sphinx:56},filenames:["about.rst","api.rst","contributing.rst","examples.rst","index.rst","installation.rst","new_bandit.rst","quick.rst"],objects:{"mabwiser.base_mab":[[1,1,1,"","BaseMAB"]],"mabwiser.base_mab.BaseMAB":[[1,2,1,"","add_arm"],[1,3,1,"","arm_to_expectation"],[1,3,1,"","arms"],[1,3,1,"","backend"],[1,3,1,"","cold_arm_to_warm_arm"],[1,2,1,"","fit"],[1,3,1,"","n_jobs"],[1,2,1,"","partial_fit"],[1,2,1,"","predict"],[1,2,1,"","predict_expectations"],[1,2,1,"","remove_arm"],[1,3,1,"","rng"],[1,3,1,"","trained_arms"],[1,2,1,"","warm_start"]],"mabwiser.mab":[[1,1,1,"","LearningPolicy"],[1,1,1,"","MAB"],[1,1,1,"","NeighborhoodPolicy"]],"mabwiser.mab.LearningPolicy":[[1,1,1,"","EpsilonGreedy"],[1,1,1,"","LinGreedy"],[1,1,1,"","LinTS"],[1,1,1,"","LinUCB"],[1,1,1,"","Popularity"],[1,1,1,"","Random"],[1,1,1,"","Softmax"],[1,1,1,"","ThompsonSampling"],[1,1,1,"","UCB1"]],"mabwiser.mab.LearningPolicy.EpsilonGreedy":[[1,3,1,"id0","epsilon"]],"mabwiser.mab.LearningPolicy.LinGreedy":[[1,3,1,"id1","epsilon"],[1,3,1,"id2","l2_lambda"],[1,3,1,"id3","scale"]],"mabwiser.mab.LearningPolicy.LinTS":[[1,3,1,"id4","alpha"],[1,3,1,"id5","l2_lambda"],[1,3,1,"id6","scale"]],"mabwiser.mab.LearningPolicy.LinUCB":[[1,3,1,"id7","alpha"],[1,3,1,"id8","l2_lambda"],[1,3,1,"id9","scale"]],"mabwiser.mab.LearningPolicy.Softmax":[[1,3,1,"id10","tau"]],"mabwiser.mab.LearningPolicy.ThompsonSampling":[[1,3,1,"id11","binarizer"]],"mabwiser.mab.LearningPolicy.UCB1":[[1,3,1,"id12","alpha"]],"mabwiser.mab.MAB":[[1,2,1,"","add_arm"],[1,3,1,"","arms"],[1,3,1,"","backend"],[1,2,1,"","fit"],[1,3,1,"","is_contextual"],[1,4,1,"id13","learning_policy"],[1,3,1,"","n_jobs"],[1,4,1,"id14","neighborhood_policy"],[1,2,1,"","partial_fit"],[1,2,1,"","predict"],[1,2,1,"","predict_expectations"],[1,2,1,"","remove_arm"],[1,3,1,"","seed"],[1,2,1,"","warm_start"]],"mabwiser.mab.NeighborhoodPolicy":[[1,1,1,"","Clusters"],[1,1,1,"","KNearest"],[1,1,1,"","LSHNearest"],[1,1,1,"","Radius"],[1,1,1,"","TreeBandit"]],"mabwiser.mab.NeighborhoodPolicy.Clusters":[[1,3,1,"id15","is_minibatch"],[1,3,1,"id16","n_clusters"]],"mabwiser.mab.NeighborhoodPolicy.KNearest":[[1,3,1,"id17","k"],[1,3,1,"id18","metric"]],"mabwiser.mab.NeighborhoodPolicy.LSHNearest":[[1,3,1,"id19","n_dimensions"],[1,3,1,"id20","n_tables"],[1,3,1,"id21","no_nhood_prob_of_arm"]],"mabwiser.mab.NeighborhoodPolicy.Radius":[[1,3,1,"id22","metric"],[1,3,1,"id23","no_nhood_prob_of_arm"],[1,3,1,"id24","radius"]],"mabwiser.mab.NeighborhoodPolicy.TreeBandit":[[1,3,1,"id27","tree_parameters"]],"mabwiser.simulator":[[1,1,1,"","Simulator"],[1,5,1,"","default_evaluator"]],"mabwiser.simulator.Simulator":[[1,3,1,"","arm_to_stats_test"],[1,3,1,"","arm_to_stats_total"],[1,3,1,"","arm_to_stats_train"],[1,3,1,"","arms"],[1,3,1,"","bandit_to_arm_to_stats_avg"],[1,3,1,"","bandit_to_arm_to_stats_max"],[1,3,1,"","bandit_to_arm_to_stats_min"],[1,3,1,"","bandit_to_arm_to_stats_neighborhoods"],[1,3,1,"","bandit_to_confusion_matrices"],[1,3,1,"","bandit_to_expectations"],[1,3,1,"","bandit_to_neighborhood_size"],[1,3,1,"","bandit_to_predictions"],[1,3,1,"","bandits"],[1,3,1,"","batch_size"],[1,3,1,"","contexts"],[1,3,1,"","decisions"],[1,3,1,"","evaluator"],[1,2,1,"","get_arm_stats"],[1,2,1,"","get_stats"],[1,3,1,"","is_ordered"],[1,3,1,"","is_quick"],[1,3,1,"","logger"],[1,2,1,"","plot"],[1,3,1,"","rewards"],[1,2,1,"","run"],[1,3,1,"","scaler"],[1,3,1,"","test_indices"],[1,3,1,"","test_size"]],"mabwiser.utils":[[1,6,1,"","Arm"],[1,1,1,"","Constants"],[1,6,1,"","Num"],[1,5,1,"","argmax"],[1,5,1,"","argmin"],[1,5,1,"","check_false"],[1,5,1,"","check_true"],[1,5,1,"","create_rng"],[1,5,1,"","reset"]],"mabwiser.utils.Constants":[[1,3,1,"","default_seed"],[1,3,1,"","distance_metrics"]],mabwiser:[[1,0,0,"-","base_mab"],[1,0,0,"-","mab"],[1,0,0,"-","simulator"],[1,0,0,"-","utils"]]},objnames:{"0":["py","module","Python module"],"1":["py","class","Python class"],"2":["py","method","Python method"],"3":["py","attribute","Python attribute"],"4":["py","property","Python property"],"5":["py","function","Python function"],"6":["py","data","Python data"]},objtypes:{"0":"py:module","1":"py:class","2":"py:method","3":"py:attribute","4":"py:property","5":"py:function","6":"py:data"},terms:{"0":[1,3,6],"00129":4,"01":3,"05":1,"1":[0,1,3,4,5,7],"10":[1,3,4],"100":3,"1000":3,"100000":3,"11":[1,3],"1109":4,"111":3,"1142":4,"12":3,"123456":[1,3,6],"13":3,"15":[1,3],"17":[1,3,4,6,7],"18":3,"19":[3,4],"1d":3,"2":[0,1,3,4,5],"20":[1,3,4,6,7],"200":3,"2010":0,"2016":0,"2017":1,"2019":4,"2021":4,"21":3,"2150021":4,"22":3,"23":3,"25":[1,3,4,6,7],"26":3,"27":3,"2d":3,"3":[1,3,4,5],"30":[1,4],"31st":4,"33":3,"37":3,"38":3,"39":3,"4":[1,3,4],"42":[3,6],"48":3,"5":[1,3],"50":3,"52":3,"55":3,"56":3,"57":3,"6":[3,4,5],"65":3,"68":3,"7":[1,3],"75":3,"8":[2,3,4],"83":3,"9":[1,3,4,6,7],"909":4,"914":4,"99":3,"abstract":[1,6],"boolean":1,"break":6,"case":[1,3,6],"class":[1,6],"default":[1,2,6],"do":[0,1,2,6],"final":6,"float":[1,6],"function":[1,2,3,6],"import":[0,1,3,4,5,6,7],"int":[1,4,6],"long":0,"new":[0,1,2,3,4],"public":[2,4],"return":[1,6],"static":1,"super":6,"true":[1,3,6],"try":5,A:[0,1,3,4,6],As:[0,1],At:0,For:[0,1,5,6],If:[1,5,6],In:[0,1,2,3,6],It:[1,3,4,6],No:1,Not:1,OR:4,Or:6,That:6,The:[0,1,2,3,4,5,6],Then:1,There:[0,3,5],These:[1,6],To:[1,5,6],_:[3,6],__init__:[1,6],__version__:6,_arm_:1,_baserng:[1,6],_copy_arm:6,_fit_arm:6,_mycoolpolici:6,_onlin:1,_parallel_fit:6,_predict_context:6,_uptake_new_arm:[1,6],_valid:6,_validate_mab_arg:6,abl:[0,6],about:[1,3,4,6],abov:6,accept:1,access:[1,6],accident:6,accompani:2,accord:6,accordingli:1,action:[0,6],activ:6,actual:1,ad:[1,3,4],adam:1,add:[1,6],add_arm:[1,3,6],addit:[1,3,6],additional_layout:3,additional_revenu:3,adher:[2,4],adjust:1,advantag:6,advertis:[0,3],affect:2,after:[0,1],ag:3,agent:0,aggreg:1,al:0,algorithm:[0,1,4],alia:1,all:[1,2,5,6],allow:[3,6],alpha:[1,3,4,7],alreadi:1,also:[0,1],altern:5,among:[0,3],amount:3,an:[0,1,3,4,6,7],ani:[1,2,5,6],api:[2,4],append:3,appli:[1,6],applic:[0,2],approach:1,approxim:1,ar:[0,1,2,3,5,6],arang:3,argmax:1,argmin:1,argument:1,arm1:[1,4,6,7],arm2:[1,4,6,7],arm3:1,arm:[1,3,6,7],arm_to_expect:[1,6],arm_to_featur:[1,3,6],arm_to_predict:6,arm_to_scal:[],arm_to_stat:1,arm_to_stats_test:1,arm_to_stats_tot:1,arm_to_stats_train:1,arm_to_threshold:1,arms_to_stat:1,arms_to_stats_train:1,arrai:[1,3,6],articl:4,artif:4,artifici:4,arxiv:0,asctim:1,assert:[3,6],assertdictequ:6,assertequ:6,assertionerror:1,assess:6,assign:1,associ:1,assum:1,assumpt:1,attribut:6,author:4,automat:6,avail:[1,3,6],averag:[1,3],avg:[1,3],b:[3,4,6],back:6,backend:[1,6],bandit:[1,3],bandit_to_arm_to_stat:1,bandit_to_arm_to_stats_avg:[1,3],bandit_to_arm_to_stats_max:[1,3],bandit_to_arm_to_stats_min:[1,3],bandit_to_arm_to_stats_neighborhood:1,bandit_to_confusion_matric:1,bandit_to_expect:1,bandit_to_neighborhood_s:1,bandit_to_predict:1,base:[1,3,4,6,7],base_mab:4,basemab:[1,6],basetest:6,basic:[1,2],batch:[1,3],batch_siz:[1,3],bdist_wheel:5,becom:3,been:[1,6],befor:[1,6],behav:6,behavior:[1,6],behind:[3,6],being:1,below:3,bernard:4,best:[0,1,3,6],best_arm:6,beta:1,between:[0,1,4,6,7],binar:[1,6],binari:1,booktitl:4,bool:[1,6],both:4,bound1:1,bound:[1,3,4],branch:6,braycurti:1,build:5,built:[3,4],cach:5,calcul:[1,6],call:[0,1,6],callabl:[1,6],can:[0,1,3,4,5,6],canberra:1,cannot:[1,6],capabl:3,captur:1,carnegi:0,cd:5,cdist:1,center:4,certain:1,chang:6,changelog:6,chebyshev:1,check:[1,2,6],check_fals:1,check_tru:1,checkout:6,choic:3,choleski:1,choos:[0,1,7],chosen:1,chronolog:1,cite:4,cityblock:1,clean:2,click:3,click_rat:3,clinic:0,clone:5,closest:1,cluster:[1,4],code:[1,2,6],coeffici:1,coher:2,cold:[1,6],cold_arm:6,cold_arm_to_warm_arm:[1,6],collect:1,collis:1,column:1,com:5,commerc:3,common:1,commun:1,compar:1,comparison:4,compat:1,complet:[1,6],compon:4,computation:1,conceptu:3,concurr:1,conf:4,confer:4,confid:[1,3,4],confirm:5,confus:1,congratul:6,connect:6,conserv:1,consid:6,consist:6,constant:1,constructor:[1,6],contain:[1,2],context:[0,1,4,6],contexts_test:3,contexts_train:3,contextu:[0,1],contextual_mab:1,continu:6,contribut:[4,6],control:1,convert:1,cool:6,copi:6,core:[1,3,6],correct:1,correctli:6,correl:1,correspond:[1,3,6],cosin:1,count:1,counter:6,cours:0,covari:1,cover:6,cpu:1,creat:[1,6],create_rng:1,cumul:[0,1],custom:[1,3,5],d:1,danc:6,data:[1,3,4,6,7],datafram:[1,3],dataset:3,dblp:4,decid:[0,6],decis:[0,1,3,4,6,7],decisions_test:3,decisions_train:3,decisiontreeclassifi:1,declar:[1,6],decomposit:1,decor:6,decreas:1,deem:1,deepcopi:6,def:[1,6],default_evalu:1,default_se:1,defin:[1,6],definit:[0,1],degre:1,denot:[1,6],depend:[1,3],descript:[1,6],deserv:6,design:[2,3],detail:6,determin:[1,6],determinist:[0,1],develop:6,dice:1,dict:[1,6],dictat:6,dictionari:[1,6],differ:[0,1,3,6],dimens:1,dimension:1,dir:5,directli:[1,6],directori:2,discov:5,displai:3,dist:5,distanc:[1,6],distance_metr:1,distance_quantil:[1,3,6],distribut:[0,1],divis:1,doc:6,docsrc:6,doctest:6,document:[2,6],doe:[1,6],doi:4,done:6,down:6,dure:[3,6],dxd:1,e:[1,3,4,6],each:[0,1,3,6],easiest:2,effect:[1,6],eg:1,elmachtoub:1,emerg:6,emili:4,empti:1,enabl:6,engin:0,ensur:[1,2],entir:1,entri:[1,6],epsilon:[1,3,4],epsilongreedi:[1,3],equal:3,error:6,estim:1,et:0,etc:1,euclidean:1,evalu:1,even:6,event:3,everi:[0,1,3,6],everyon:6,evolv:6,exampl:[1,4,5,6,7],excel:[0,4],except:1,exchang:1,execut:[1,6],exist:[0,1,2,6],exp:6,expect:[0,1,3,4,6,7],expens:1,experi:[1,3],explor:[1,3],expos:4,express:[1,6],extend:[1,6],face:0,factor:1,fals:[1,3,6],fashion:6,fast:[1,4],featur:[1,2,3,4,6],fidel:[4,5],field:[1,3,6],file:6,filter:1,first:[1,6],fit:[1,3,4,6,7],fit_arm:6,fit_transform:3,flag:[1,6],folder:[3,5,6],follow:[1,4,6],foremost:6,format:1,formula:6,frac:1,frame:[1,3],framework:6,free:[1,4],from:[0,1,3,4,5,6,7],fromkei:6,g:[1,6],gener:[1,2,3,6],get:[1,2],get_arm_stat:1,get_stat:1,git:[5,6],github:[2,4,5,6],give:1,given:[0,1,3,6],global:1,go:6,goe:[1,6],good:1,greater:1,greedi:[1,3,4,6],greedili:1,guarante:1,guidelin:6,ha:[0,1,6],ham:1,handl:6,happi:2,hash:1,have:[0,1,2,6],heavili:4,helper:1,here:[3,6],high:3,higher:1,highest:[0,1],hint:6,histor:[1,3],histori:[1,3],homepag:3,host:[2,4],how:[1,3,4,5,6,7],http:[4,5],hyper:[1,3,4,6],hyper_parameter_tun:3,hyperplan:1,i:[1,3],i_d:1,ictai:4,idea:[3,6],ident:1,ieee:4,ijait:4,imagin:6,immut:6,implement:[1,4],implementor:[1,6],includ:[5,6],incorpor:1,increas:1,increment:6,independ:3,index:[1,4],indic:1,induc:1,infin:1,inform:[0,3],inherit:[1,6],initi:[1,6],inner:6,inproceed:4,input:[1,6],instal:[4,6],instead:[1,6],integ:1,intel:4,intellig:4,intend:1,interfac:[1,6],intern:[1,4,6],internet:0,interpret:1,introduc:[1,6],invers:1,invest:4,involv:6,is_contextu:1,is_minibatch:1,is_ord:[1,3],is_per_arm:[1,3],is_predict:6,is_quick:1,isn:6,issu:[2,4],item:[1,6],its:[1,6],j:4,jaccard:1,joblib:1,journal:4,k:[0,1,4],kadioglu:4,keep:1,kei:1,keyword:6,kleynhan:4,kmean:1,knearest:1,know:0,kulsinski:1,kwarg:1,l2_lambda:[1,3],lambda:1,larg:3,later:6,latest:5,layout:3,lead:3,leaf:1,learn:[0,1,3,4,6,7],learning_polici:[1,3,6],learningpolici:[1,3,4,5,6,7],learnt:3,least:1,leav:1,length:1,less:1,let:6,level:3,levelnam:1,leverag:6,li:0,librari:[1,4,6,7],like:6,likelihood:1,linear:3,lingreedi:[1,4],lint:[1,4],linucb:[1,3,4],list:[1,3,5,6],list_of_arm:1,listen:3,live:6,local:[1,6],lock:1,log:1,log_fil:1,log_format:1,logger:1,logist:1,loki:1,look:6,lot:1,low:1,lp:6,lsh:4,lshnearest:1,lu:0,m:5,mab1:1,mab2:1,mab:[0,4,5,6,7],mab_nam:3,mabwis:[2,3,5,6,7],machin:0,made:1,magic:6,mahalanobi:1,mai:0,main:[2,6],make:[0,1,2,3,4,6],make_classif:3,mani:[0,1],map:[1,6],marek:1,master:[5,6],match:1,math:6,matric:1,matrix:1,max:1,maxim:0,maximum:1,mcnelli:1,md:6,mean:[1,3],mean_:[],mellon:0,member:6,memori:1,merg:6,messag:1,meta:6,method:[1,2,3,6],metric:[1,3],might:6,min:1,minibatchkmean:1,minimum:[1,6],minkowski:1,mismatch:1,miss:1,mode:6,model:[0,1,3,4,6,7],model_select:3,model_to_confusion_matric:1,model_to_predict:1,models_to_evalu:1,modif:2,modifi:1,modul:[1,4],more:[0,1],most:6,move:6,mu:1,mu_i:1,much:0,multi:1,multipl:[0,1,3],multipli:1,multiprocess:1,multivari:1,music:3,must:1,my_paramet:6,my_results_:3,my_value_to_arm:6,mycoolbandittest:6,mycoolpolici:6,n:[1,3],n_cluster:1,n_context_col:1,n_dimens:1,n_featur:3,n_i:1,n_inform:3,n_job:[1,3,6],n_row:1,n_sampl:3,n_tabl:1,name:[1,6],namedtupl:6,nan:1,nd:1,ndarrai:[1,6],nearest:[1,4],necessari:[5,6],need:[0,3,6],neg:[1,6],neighbor:1,neighborhood:[1,3,4,6],neighborhood_polici:[1,3],neighborhoodpolici:[1,3,4,5,6,7],next:[3,6],nn:1,no_nhood_prob_of_arm:1,non:[1,4],none:[1,3,5,6],noreturn:[1,6],normal:1,note:[1,6],notebook:5,notic:[1,6],notimplementederror:1,novemb:4,now:6,np:[1,3,6],nparrai:1,num:[1,6],num_run:6,number:[1,4,6],numer:1,numpi:[1,3,6],numpydoc:2,o:1,object:[0,1,6],observ:[0,1,3,6],obtain:1,offlin:[1,3],offline_sim:1,oh:1,older:2,one:[1,6],onli:[0,1],onlin:[0,1,3,6],onto:[1,6],oper:[1,6],optim:3,option:[0,1,3,6],order:[1,3],org:4,origin:5,other:[0,1,3,6],otherwis:1,our:[2,6],out:[1,2],outcom:0,output:1,over:1,overal:0,overhead:1,overrid:1,own:6,p:1,packag:[5,6],page:4,pairwis:6,panda:[1,3],parallel:[1,4,6],paralleliz:4,param:6,paramet:[1,3,4,6],parametr:[1,4],part:[0,6],partial:1,partial_fit:[1,3,6],partit:1,pass:[2,5,6],past:3,pd:[1,3],pep:[2,4],per:3,per_arm:1,perform:[1,6],petrik:1,phase:[3,6],pip:[5,6],pipelin:1,plai:0,platform:5,playlist:3,pleas:[4,6],plot:[1,3],point:[1,6],polici:[1,3,6,7],pool:1,popular:[1,4,6],portland:4,posit:1,possibl:[3,6],potenti:[0,1],practic:1,precis:1,predict:[1,3,4,6,7],predict_expect:[1,3,6],prefix:[3,6],prepar:6,preprocess:[1,3],present:1,previou:[1,2,3,6],print:3,privat:6,probabl:1,probe:3,problem:[0,1,3],proc:0,procedur:1,process:1,progress:1,project:[1,2,5],promot:1,properli:6,properti:[1,6],proportion:1,prototyp:[1,4],provid:[1,4,6],publish:4,pull:[0,4,5],py3:5,py:[5,6],pydoc:6,python:[1,4,5,6],quantil:1,question:0,quickli:6,radiu:[1,3,4],rais:[1,6],randint:3,random:[1,3,4,6],random_st:3,randomli:1,rang:3,rate:3,rather:1,ration:1,re:3,read:1,readi:6,readm:6,real:[0,6],realiti:3,receiv:[0,1],recommend:[0,1,3],recompil:6,refer:[0,5],reflect:1,region:1,regist:6,regress:1,regular:1,releas:4,reli:1,remain:6,remov:[1,6],remove_arm:1,renew:0,repeat:1,repo:[2,3,5],represent:1,request:4,requir:[1,2,4],research:[0,1,4],reset:[1,6],resourc:0,result:[0,1,3],retriev:[1,3],revenu:3,review:6,reward:[0,1,3,4,6,7],rewards_test:3,rewards_train:3,ridg:1,rng:[1,6],robust:1,rogerstanimoto:1,root:1,routin:1,row:1,rst:6,run:[0,1,3,5,6],russellrao:1,ryan:1,s0218213021500214:4,s:[1,3,4,6],same:[1,6],sampl:[0,1,2,4],save:3,save_result:3,scale:[1,3],scaler:[1,3,6],scenario:3,scikit:4,scipi:1,scratch:[5,6],sdist:5,search:6,sechan:1,section:[2,6],see:[1,6],seed:[1,3,6],seen:6,select:1,self:6,send:4,sensit:1,separ:1,sequenc:0,sequenti:0,serdar:4,seri:[1,3,6],servic:3,set:[0,1,3,6],setup:4,setuptool:5,seuclidean:1,share:[1,6],shell:5,should:[1,2,3,5],show:[3,4,7],shown:[3,6],side:0,sigma:1,sign:1,signatur:[1,6],sim:3,simhash:1,similar:6,simpl:[1,6],simul:4,sinc:3,singl:1,situat:0,size:[1,3],skeleton:1,skip:[1,3],sklearn:[1,3],small:6,so:[1,6],softmax:[1,4],sokalmichen:1,sokalsneath:1,solv:[1,3],some:[0,1,6],someth:0,space:1,spatial:1,specif:1,specifi:1,speed:1,sphinx:6,split:[1,3],sqeuclidean:1,sqrt:1,squar:1,standard:[2,3,4],standardscal:[1,3],start:[1,2,3,6],start_index:[1,6],stat:[1,3],statis:3,statist:1,std:1,step:[1,3,6],stick:0,still:2,stochast:0,store:[1,6],str:[1,3,6],stream:3,strength:1,strengthen:6,string:[1,6],strong:4,strongkk19:4,strongkk21:4,style:4,sub:1,subscrib:3,success:[1,5,6],suffer:1,suit:6,suitabl:0,sum:[1,6],support:[1,3,4],sure:[2,6],survei:0,system:6,t:[0,1,6],tabl:1,take:[1,6],tau:1,tell:6,temperatur:1,test:[0,1,2,3,4,7],test_:6,test_add_new_arm:6,test_bas:6,test_df:3,test_df_revenu:3,test_fit_twic:6,test_indic:1,test_input_typ:6,test_invalid:6,test_my_paramet:6,test_mycoolbandit:6,test_parallel:6,test_partial_fit:6,test_remove_existing_arm:6,test_simple_usecase_arm:6,test_simple_usecase_expect:6,test_siz:[1,3],test_unused_arm:6,test_warm_start:6,test_within_neighborhood_polici:6,test_zero_reward:6,than:1,thei:[0,1,6],them:1,thi:[0,1,3,6],thompson:[1,4],thompsonsampl:1,those:6,thread:1,three:6,time:[0,1,6],titl:4,todo:6,togeth:1,toi:3,too:6,tool:4,total:[1,3,6],track:[1,4],train:[1,3,4,6,7],train_df:3,train_test_split:[1,3],trained_arm:1,transform:3,tree:1,tree_paramet:1,treebandit:[1,4],treeheurist:1,trial:[0,1],ts:4,tune:[1,3,4],tupl:1,two:[4,5,6,7],txt:[5,6],tyler:0,type:[1,6],typeerror:1,typic:6,uai:1,ucb1:[1,3,4,7],ucb:1,uncertainti:0,under:[0,1,6],underli:[0,6],uniform:1,uniformli:1,uninstal:6,union:1,uniqu:1,unit:[1,2,6],unittest:[5,6],univers:0,unknown:0,unlik:6,untrain:[1,6],unus:6,up:1,updat:[1,3,6],upgrad:4,upper:[1,3,4],url:4,us:[1,2,3,4,5,6,7],usa:4,usag:[4,5,6],user:[3,6],util:[4,6],valid:[1,6],valu:[1,6],valueerror:1,var_:[],varianc:1,variant:0,vector:1,veri:1,version:[1,5,6],via:6,victori:6,violat:1,volum:4,wa:[1,5],wai:2,want:[2,3],warm:[1,3,6],warm_arm:6,warm_start:[1,3,6],we:[0,2,3,6],websit:3,weight:[1,3],were:6,what:[0,1,3,6],wheel:5,when:[1,6],where:[0,1,6],whether:[1,3],which:[0,1,3,6],whl:5,within:[1,6],work:[1,2,6],worker:1,world:[0,6],worst:3,would:6,wrapper:6,x:[1,5],x_i:1,year:4,yet:0,you:[1,2,3,4,5,6],your:[4,6],zero:[1,6],zhou:0},titles:["About Multi-Armed Bandits","MABWiser Public API","Contributing","Usage Examples","MABWiser Contextual Multi-Armed Bandits","Installation","Adding a New Bandit","Quick Start"],titleterms:{"1":6,"2":6,"3":6,"4":6,"new":6,"public":[1,6],about:0,ad:6,algorithm:6,api:[1,6],arm:[0,4],avail:4,bandit:[0,4,6],base_mab:1,bug:4,citat:4,code:[4,5],content:4,context:3,contextu:[3,4],contribut:2,exampl:3,exploit:0,explor:0,expos:6,free:3,high:6,implement:6,indic:4,instal:5,level:6,librari:5,mab:[1,3],mabwis:[1,4],multi:[0,4],non:3,option:5,overview:6,parallel:3,parametr:3,polici:4,pull:6,quick:[4,7],report:4,request:6,requir:5,send:6,setup:5,simul:[1,3],sourc:[4,5],start:[4,7],tabl:4,test:[5,6],upgrad:5,usag:3,util:1,vs:0,your:5}})
\ No newline at end of file
diff --git a/mabwiser/_version.py b/mabwiser/_version.py
index 1df14e0..37c0bb0 100644
--- a/mabwiser/_version.py
+++ b/mabwiser/_version.py
@@ -3,5 +3,5 @@
 
 __author__ = "FMR LLC"
 __email__ = "opensource@fmr.com"
-__version__ = "2.3.0"
+__version__ = "2.4.0"
 __copyright__ = "Copyright (C), FMR LLC"
diff --git a/mabwiser/base_mab.py b/mabwiser/base_mab.py
index c024095..c3294a1 100755
--- a/mabwiser/base_mab.py
+++ b/mabwiser/base_mab.py
@@ -7,7 +7,7 @@
 
 import abc
 from itertools import chain
-from typing import Callable, Dict, List, NoReturn, Optional
+from typing import Callable, Dict, List, NoReturn, Optional, Union
 import multiprocessing as mp
 
 from joblib import Parallel, delayed
@@ -57,7 +57,7 @@ class BaseMAB(metaclass=abc.ABCMeta):
         If set to -2, all CPUs but one are used, and so on.
     backend: str, optional
         Specify a parallelization backend implementation supported in the joblib library. Supported options are:
-        - “loky” used by default, can induce some communication and memory overhead when exchanging input and output data with the worker Python processes.
+        - “loky” used by default, can induce some communication and memory overhead when exchanging input and output.
         - “multiprocessing” previous process-based backend based on multiprocessing.Pool. Less robust than loky.
         - “threading” is a very low-overhead backend but it suffers from the Python Global Interpreter Lock if the
           called function relies a lot on Python objects.
@@ -86,14 +86,14 @@ def __init__(self, rng: _BaseRNG, arms: List[Arm], n_jobs: int, backend: str = N
         self.cold_arm_to_warm_arm: Dict[Arm, Arm] = dict()
         self.trained_arms: List[Arm] = list()
 
-    def add_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None) -> NoReturn:
+    def add_arm(self, arm: Arm, binarizer: Callable = None) -> NoReturn:
         """Introduces a new arm to the bandit.
 
         Adds the new arm with zero expectations and
         calls the ``_uptake_new_arm()`` function of the sub-class.
         """
         self.arm_to_expectation[arm] = 0
-        self._uptake_new_arm(arm, binarizer, scaler)
+        self._uptake_new_arm(arm, binarizer)
 
     def remove_arm(self, arm: Arm) -> NoReturn:
         """Removes arm from the bandit.
@@ -122,7 +122,7 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray,
         pass
 
     @abc.abstractmethod
-    def predict(self, contexts: Optional[np.ndarray] = None) -> Arm:
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
         """Abstract method.
 
         Returns the predicted arm.
@@ -130,7 +130,8 @@ def predict(self, contexts: Optional[np.ndarray] = None) -> Arm:
         pass
 
     @abc.abstractmethod
-    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
         """Abstract method.
 
         Returns a dictionary from arms (keys) to their expected rewards (values).
@@ -146,7 +147,7 @@ def _copy_arms(self, cold_arm_to_warm_arm: Dict[Arm, Arm]) -> NoReturn:
         pass
 
     @abc.abstractmethod
-    def _uptake_new_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None) -> NoReturn:
+    def _uptake_new_arm(self, arm: Arm, binarizer: Callable = None) -> NoReturn:
         """Abstract method.
 
         Updates the multi-armed bandit with the new arm.
diff --git a/mabwiser/clusters.py b/mabwiser/clusters.py
index 60f0ba0..1e9089a 100755
--- a/mabwiser/clusters.py
+++ b/mabwiser/clusters.py
@@ -84,11 +84,11 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray,
 
         self._fit_operation()
 
-    def predict(self, contexts: Optional[np.ndarray] = None):
+    def predict(self, contexts: np.ndarray = None) -> Union[Arm, List[Arm]]:
         # Return predict within the cluster
         return self._parallel_predict(contexts, is_predict=True)
 
-    def predict_expectations(self, contexts: Optional[np.ndarray] = None):
+    def predict_expectations(self, contexts: np.ndarray = None) -> Union[Dict[Arm, Num], List[Dict[Arm, Num]]]:
         # Return predict expectations within the cluster
         return self._parallel_predict(contexts, is_predict=False)
 
diff --git a/mabwiser/greedy.py b/mabwiser/greedy.py
index 7f17443..2615349 100755
--- a/mabwiser/greedy.py
+++ b/mabwiser/greedy.py
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
-from typing import Callable, Dict, List, NoReturn, Optional
+from typing import Callable, Dict, List, NoReturn, Optional, Union
 
 import numpy as np
 
@@ -32,19 +32,33 @@ def fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.ndarray =
     def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.ndarray = None) -> NoReturn:
         self._parallel_fit(decisions, rewards, contexts)
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
 
-        # Return the first arm with maximum expectation
-        return argmax(self.predict_expectations())
+        # Return the arm with maximum expectation
+        expectations = self.predict_expectations(contexts)
+        if isinstance(expectations, dict):
+            return argmax(expectations)
+        else:
+            return [argmax(exp) for exp in expectations]
 
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
 
         # Return a random expectation (between 0 and 1) for each arm with epsilon probability,
-        # and the actual arm expectations otherwise
-        if self.rng.rand() < self.epsilon:
-            return dict((arm, self.rng.rand()) for arm in self.arms).copy()
+        # and the actual arm expectations otherwise.
+        # If contexts is None or has length of 1 generate single arm to expectations,
+        # otherwise use vectorized functions to generate a list of arm to expectations with same length as contexts.
+        if contexts is None or len(contexts) == 1:
+            if self.rng.rand() < self.epsilon:
+                return dict((arm, self.rng.rand()) for arm in self.arms).copy()
+            else:
+                return self.arm_to_expectation.copy()
         else:
-            return self.arm_to_expectation.copy()
+            probability = self.rng.rand(len(contexts))
+            random_values = self.rng.rand((len(contexts), len(self.arms)))
+            expectations = [dict(zip(self.arms, exp)).copy() if probability[index] < self.epsilon
+                            else self.arm_to_expectation.copy() for index, exp in enumerate(random_values)]
+            return expectations
 
     def _copy_arms(self, cold_arm_to_warm_arm):
         for cold_arm, warm_arm in cold_arm_to_warm_arm.items():
diff --git a/mabwiser/linear.py b/mabwiser/linear.py
index 6816c38..d98800d 100755
--- a/mabwiser/linear.py
+++ b/mabwiser/linear.py
@@ -2,9 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
-from typing import Callable, Dict, List, NoReturn, Optional
+from typing import Callable, Dict, List, NoReturn, Optional, Union
 
 import numpy as np
+from sklearn.preprocessing import StandardScaler
 
 from mabwiser.base_mab import BaseMAB
 from mabwiser.utils import Arm, Num, argmax, _BaseRNG, create_rng
@@ -12,19 +13,19 @@
 
 class _RidgeRegression:
 
-    def __init__(self, rng: _BaseRNG, alpha: Num = 1.0, l2_lambda: Num = 1.0, 
-                 scaler: Optional[Callable] = None):
+    def __init__(self, rng: _BaseRNG, alpha: Num = 1.0, l2_lambda: Num = 1.0, scale: bool = False):
 
         # Ridge Regression: https://onlinecourses.science.psu.edu/stat857/node/155/
         self.rng = rng                      # random number generator
         self.alpha = alpha                  # exploration parameter
         self.l2_lambda = l2_lambda          # regularization parameter
-        self.scaler = scaler                # standard scaler object
+        self.scale = scale                  # scale contexts
 
         self.beta = None                    # (XtX + l2_lambda * I_d)^-1 * Xty = A^-1 * Xty
         self.A = None                       # (XtX + l2_lambda * I_d)
         self.A_inv = None                   # (XtX + l2_lambda * I_d)^-1
         self.Xty = None
+        self.scaler = None
 
     def init(self, num_features):
         # By default, assume that
@@ -33,12 +34,18 @@ def init(self, num_features):
         self.A = self.l2_lambda * np.identity(num_features)
         self.A_inv = self.A.copy()
         self.beta = np.dot(self.A_inv, self.Xty)
+        self.scaler = StandardScaler() if self.scale else None
 
     def fit(self, X, y):
 
         # Scale
         if self.scaler is not None:
-            X = self.scaler.transform(X.astype('float64'))
+            X = X.astype('float64')
+            if not hasattr(self.scaler, 'scale_'):
+                self.scaler.fit(X)
+            else:
+                self.scaler.partial_fit(X)
+            X = self.scaler.transform(X)
 
         # X transpose
         Xt = X.T
@@ -63,6 +70,9 @@ def predict(self, x):
         return np.dot(x, self.beta)
 
     def _scale_predict_context(self, x):
+        if not hasattr(self.scaler, 'scale_'):
+            return x
+
         # Reshape 1D array to 2D
         x = x.reshape(1, -1)
 
@@ -106,22 +116,17 @@ class _Linear(BaseMAB):
     factory = {"ts": _LinTS, "ucb": _LinUCB, "ridge": _RidgeRegression}
 
     def __init__(self, rng: _BaseRNG, arms: List[Arm], n_jobs: int, backend: Optional[str],
-                 alpha: Num, epsilon: Num, l2_lambda: Num, regression: str,
-                 arm_to_scaler: Optional[Dict[Arm, Callable]] = None):
+                 alpha: Num, epsilon: Num, l2_lambda: Num, regression: str, scale: bool):
         super().__init__(rng, arms, n_jobs, backend)
         self.alpha = alpha
         self.epsilon = epsilon
         self.l2_lambda = l2_lambda
         self.regression = regression
-
-        # Create ridge regression model for each arm
+        self.scale = scale
         self.num_features = None
 
-        if arm_to_scaler is None:
-            arm_to_scaler = dict((arm, None) for arm in arms)
-
-        self.arm_to_model = dict((arm, _Linear.factory.get(regression)(rng, alpha, 
-                                                                       l2_lambda, arm_to_scaler[arm])) for arm in arms)
+        # Create regression model for each arm
+        self.arm_to_model = dict((arm, _Linear.factory.get(regression)(rng, alpha, l2_lambda, scale)) for arm in arms)
 
     def fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.ndarray = None) -> NoReturn:
 
@@ -137,11 +142,11 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.n
         # Perform parallel fit
         self._parallel_fit(decisions, rewards, contexts)
 
-    def predict(self, contexts: np.ndarray = None):
+    def predict(self, contexts: np.ndarray = None) -> Union[Arm, List[Arm]]:
         # Return predict for the given context
         return self._parallel_predict(contexts, is_predict=True)
 
-    def predict_expectations(self, contexts: np.ndarray = None):
+    def predict_expectations(self, contexts: np.ndarray = None) -> Union[Dict[Arm, Num], List[Dict[Arm, Num]]]:
         # Return predict expectations for the given context
         return self._parallel_predict(contexts, is_predict=False)
 
@@ -149,10 +154,10 @@ def _copy_arms(self, cold_arm_to_warm_arm):
         for cold_arm, warm_arm in cold_arm_to_warm_arm.items():
             self.arm_to_model[cold_arm] = deepcopy(self.arm_to_model[warm_arm])
 
-    def _uptake_new_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None):
+    def _uptake_new_arm(self, arm: Arm, binarizer: Callable = None):
 
         # Add to untrained_arms arms
-        self.arm_to_model[arm] = _Linear.factory.get(self.regression)(self.rng, self.alpha, self.l2_lambda, scaler)
+        self.arm_to_model[arm] = _Linear.factory.get(self.regression)(self.rng, self.alpha, self.l2_lambda, self.scale)
 
         # If fit happened, initialize the new arm to defaults
         is_fitted = self.num_features is not None
diff --git a/mabwiser/mab.py b/mabwiser/mab.py
index 4ba14a0..c516c71 100755
--- a/mabwiser/mab.py
+++ b/mabwiser/mab.py
@@ -86,12 +86,10 @@ class LinGreedy(NamedTuple):
             The regularization strength.
             Integer or float. Cannot be negative.
             Default value is 1.0.
-        arm_to_scaler: Dict[Arm, Callable]
-            Standardize context features by arm.
-            Dictionary mapping each arm to a scaler object. It is assumed
-            that the scaler objects are already fit and will only be used
-            to transform context features.
-            Default value is None.
+        scale: bool
+            Whether to scale features to have zero mean and unit variance.
+            Uses StandardScaler in sklearn.preprocessing.
+            Default value is False.
 
         Example
         -------
@@ -107,15 +105,14 @@ class LinGreedy(NamedTuple):
         """
         epsilon: Num = 0.1
         l2_lambda: Num = 1.0
-        arm_to_scaler: Dict[Arm, Callable] = None
+        scale: bool = False
 
         def _validate(self):
             check_true(isinstance(self.epsilon, (int, float)), TypeError("Epsilon must be an integer or float."))
             check_true(0 <= self.epsilon <= 1, ValueError("Epsilon must be between zero and one."))
             check_true(isinstance(self.l2_lambda, (int, float)), TypeError("L2_lambda must be an integer or float."))
             check_true(0 <= self.l2_lambda, ValueError("The value of l2_lambda cannot be negative."))
-            if self.arm_to_scaler is not None:
-                check_true(isinstance(self.arm_to_scaler, dict), TypeError("Arm_to_scaler must be a dictionary"))
+            check_true(isinstance(self.scale, bool), TypeError("Standardize must be True or False."))
 
     class LinTS(NamedTuple):
         """ LinTS Learning Policy
@@ -148,12 +145,10 @@ class LinTS(NamedTuple):
             The regularization strength.
             Integer or float. Must be greater than zero.
             Default value is 1.0.
-        arm_to_scaler: Dict[Arm, Callable]
-            Standardize context features by arm.
-            Dictionary mapping each arm to a scaler object. It is assumed
-            that the scaler objects are already fit and will only be used
-            to transform context features.
-            Default value is None.
+        scale: bool
+            Whether to scale features to have zero mean and unit variance.
+            Uses StandardScaler in sklearn.preprocessing.
+            Default value is False.
 
         Example
         -------
@@ -169,15 +164,14 @@ class LinTS(NamedTuple):
         """
         alpha: Num = 1.0
         l2_lambda: Num = 1.0
-        arm_to_scaler: Dict[Arm, Callable] = None
+        scale: bool = False
 
         def _validate(self):
             check_true(isinstance(self.alpha, (int, float)), TypeError("Alpha must be an integer or float."))
             check_true(0 < self.alpha, ValueError("The value of alpha must be greater than zero."))
             check_true(isinstance(self.l2_lambda, (int, float)), TypeError("L2_lambda must be an integer or float."))
             check_true(0 < self.l2_lambda, ValueError("The value of l2_lambda must be greater than zero."))
-            if self.arm_to_scaler is not None:
-                check_true(isinstance(self.arm_to_scaler, dict), TypeError("Arm_to_scaler must be a dictionary"))
+            check_true(isinstance(self.scale, bool), TypeError("Scale must be True or False."))
 
     class LinUCB(NamedTuple):
         """LinUCB Learning Policy.
@@ -208,12 +202,10 @@ class LinUCB(NamedTuple):
             The regularization strength.
             Integer or float. Cannot be negative.
             Default value is 1.0.
-        arm_to_scaler: Dict[Arm, Callable]
-            Standardize context features by arm.
-            Dictionary mapping each arm to a scaler object. It is assumed
-            that the scaler objects are already fit and will only be used
-            to transform context features.
-            Default value is None.
+        scale: bool
+            Whether to scale features to have zero mean and unit variance.
+            Uses StandardScaler in sklearn.preprocessing.
+            Default value is False.
 
         Example
         -------
@@ -229,15 +221,14 @@ class LinUCB(NamedTuple):
         """
         alpha: Num = 1.0
         l2_lambda: Num = 1.0
-        arm_to_scaler: Dict[Arm, Callable] = None
+        scale: bool = False
 
         def _validate(self):
             check_true(isinstance(self.alpha, (int, float)), TypeError("Alpha must be an integer or float."))
             check_true(0 <= self.alpha, ValueError("The value of alpha cannot be negative."))
             check_true(isinstance(self.l2_lambda, (int, float)), TypeError("L2_lambda must be an integer or float."))
             check_true(0 <= self.l2_lambda, ValueError("The value of l2_lambda cannot be negative."))
-            if self.arm_to_scaler is not None:
-                check_true(isinstance(self.arm_to_scaler, dict), TypeError("Arm_to_scaler must be a dictionary"))
+            check_true(isinstance(self.scale, bool), TypeError("Scale must be True or False."))
 
     class Popularity(NamedTuple):
         """Randomized Popularity Learning Policy.
@@ -299,8 +290,8 @@ class Softmax(NamedTuple):
             P(arm) = \\frac{ e ^  \\frac{\\mu_i - \\max{\\mu}}{ \\tau } }
             { \\Sigma{e ^  \\frac{\\mu - \\max{\\mu}}{ \\tau }}  }
 
-        where :math:`\\mu_i` is the mean reward for that arm and :math:`\\tau` is the "temperature" to determine the degree of
-        exploration.
+        where :math:`\\mu_i` is the mean reward for that arm and :math:`\\tau` is the "temperature" to determine
+        the degree of exploration.
 
         Attributes
         ----------
@@ -537,7 +528,7 @@ class LSHNearest(NamedTuple):
             Integer value. Must be greater than zero.
             Default value is 3.
         no_nhood_prob_of_arm: None or List
-            The probabilities associated with each arm. Used to select random arm if a prediction context has no neighbors.
+            The probabilities associated with each arm. Used to select random arm if context has no neighbors.
             If not given, a uniform random distribution over all arms is assumed.
             The probabilities should sum up to 1.
 
@@ -564,7 +555,7 @@ def _validate(self):
             check_true(self.n_dimensions > 0, ValueError("n_dimensions must be greater than zero."))
             check_true(isinstance(self.n_tables, int), TypeError("n_tables must be an integer"))
             check_true(self.n_tables > 0, ValueError("n_tables must be greater than zero."))
-            check_true((self.no_nhood_prob_of_arm == None) or isinstance(self.no_nhood_prob_of_arm, List),
+            check_true((self.no_nhood_prob_of_arm is None) or isinstance(self.no_nhood_prob_of_arm, List),
                        TypeError("no_nhood_prob_of_arm must be None or List."))
             if isinstance(self.no_nhood_prob_of_arm, List):
                 check_true(np.isclose(sum(self.no_nhood_prob_of_arm), 1.0),
@@ -587,7 +578,7 @@ class Radius(NamedTuple):
             Accepts any of the metrics supported by scipy.spatial.distance.cdist.
             Default value is Euclidean distance.
         no_nhood_prob_of_arm: None or List
-            The probabilities associated with each arm. Used to select random arm if a prediction context has no neighbors.
+            The probabilities associated with each arm. Used to select random arm if context has no neighbors.
             If not given, a uniform random distribution over all arms is assumed.
             The probabilities should sum up to 1.
 
@@ -614,7 +605,7 @@ def _validate(self):
             check_true((self.metric in Constants.distance_metrics),
                        ValueError("Metric must be supported by scipy.spatial.distance.cdist"))
             check_true(self.radius > 0, ValueError("Radius must be greater than zero."))
-            check_true((self.no_nhood_prob_of_arm == None) or isinstance(self.no_nhood_prob_of_arm, List),
+            check_true((self.no_nhood_prob_of_arm is None) or isinstance(self.no_nhood_prob_of_arm, List),
                        TypeError("no_nhood_prob_of_arm must be None or List."))
             if isinstance(self.no_nhood_prob_of_arm, List):
                 check_true(np.isclose(sum(self.no_nhood_prob_of_arm), 1.0),
@@ -868,13 +859,13 @@ def __init__(self,
             lp = _UCB1(self._rng, self.arms, self.n_jobs, self.backend, learning_policy.alpha)
         elif isinstance(learning_policy, LearningPolicy.LinGreedy):
             lp = _Linear(self._rng, self.arms, self.n_jobs, self.backend, 0, learning_policy.epsilon,
-                          learning_policy.l2_lambda, "ridge", learning_policy.arm_to_scaler)
+                         learning_policy.l2_lambda, "ridge", learning_policy.scale)
         elif isinstance(learning_policy, LearningPolicy.LinTS):
             lp = _Linear(self._rng, self.arms, self.n_jobs, self.backend, learning_policy.alpha, 0,
-                          learning_policy.l2_lambda, "ts", learning_policy.arm_to_scaler)
+                         learning_policy.l2_lambda, "ts", learning_policy.scale)
         elif isinstance(learning_policy, LearningPolicy.LinUCB):
             lp = _Linear(self._rng, self.arms, self.n_jobs, self.backend, learning_policy.alpha, 0,
-                          learning_policy.l2_lambda, "ucb", learning_policy.arm_to_scaler)
+                         learning_policy.l2_lambda, "ucb", learning_policy.scale)
         else:
             check_true(False, ValueError("Undefined learning policy " + str(learning_policy)))
 
@@ -936,15 +927,12 @@ def learning_policy(self):
             else:
                 return LearningPolicy.EpsilonGreedy(lp.epsilon)
         elif isinstance(lp, _Linear):
-            arm_to_scaler = dict()
-            for arm in lp.arms:
-                arm_to_scaler[arm] = lp.arm_to_model[arm].scaler
             if lp.regression == 'ridge':
-                return LearningPolicy.LinGreedy(lp.epsilon, lp.l2_lambda, arm_to_scaler)
+                return LearningPolicy.LinGreedy(lp.epsilon, lp.l2_lambda, lp.scale)
             elif lp.regression == 'ts':
-                return LearningPolicy.LinTS(lp.alpha, lp.l2_lambda, arm_to_scaler)
+                return LearningPolicy.LinTS(lp.alpha, lp.l2_lambda, lp.scale)
             elif lp.regression == 'ucb':
-                return LearningPolicy.LinUCB(lp.alpha, lp.l2_lambda, arm_to_scaler)
+                return LearningPolicy.LinUCB(lp.alpha, lp.l2_lambda, lp.scale)
             else:
                 check_true(False, ValueError("Undefined regression " + str(lp.regression)))
         elif isinstance(lp, _Random):
@@ -981,7 +969,7 @@ def neighborhood_policy(self):
         else:
             return None
 
-    def add_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None) -> NoReturn:
+    def add_arm(self, arm: Arm, binarizer: Callable = None) -> NoReturn:
         """ Adds an _arm_ to the list of arms.
 
         Incorporates the arm into the learning and neighborhood policies with no training data.
@@ -992,8 +980,6 @@ def add_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None)
             The new arm to be added.
         binarizer: Callable
             The new binarizer function for Thompson Sampling.
-        scaler: Callable
-            A scaler object from sklearn.preprocessing.
 
         Returns
         -------
@@ -1002,8 +988,6 @@ def add_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None)
         Raises
         ------
         TypeError:  For ThompsonSampling, binarizer must be a callable function.
-        TypeError:  The standard scaler object must have a transform method.
-        TypeError:  The standard scaler object must be fit with calculated ``mean_`` and ``var_`` attributes.
 
         ValueError: A binarizer function was provided but the learning policy is not Thompson Sampling.
         ValueError: The arm already exists.
@@ -1020,17 +1004,11 @@ def add_arm(self, arm: Arm, binarizer: Callable = None, scaler: Callable = None)
                                                       "success for a given arm decision. Specifically, the function "
                                                       "signature is binarize(arm: Arm, reward: Num) -> True/False "
                                                       "or 0/1"))
-
-        if scaler:
-            check_true(hasattr(scaler, 'transform'),
-                       TypeError("Scaler must be a scaler object from sklearn.preprocessing with a transform method"))
-            check_true(hasattr(scaler, 'mean_') and hasattr(scaler, 'var_'),
-                       TypeError("Scaler must be fit with calculated mean_ and var_ attributes"))
         check_false(arm in self.arms, ValueError("The arm is already in the list of arms."))
 
         self._validate_arm(arm)
         self.arms.append(arm)
-        self._imp.add_arm(arm, binarizer, scaler)
+        self._imp.add_arm(arm, binarizer)
 
     def remove_arm(self, arm: Arm) -> NoReturn:
         """Removes an _arm_ from the list of arms.
@@ -1184,9 +1162,10 @@ def predict(self,
 
         Parameters
         ----------
-        contexts : Union[None, List[List[Num]], np.ndarray, pd.Series, pd.DataFrame]
-            The context under which each decision is made. Default value is None.
-            Contexts should be ``None`` for context-free bandits and is required for contextual bandits.
+        contexts : Union[None, List[Num], List[List[Num]], np.ndarray, pd.Series, pd.DataFrame]
+            The context for the expected rewards. Default value is None.
+            If contexts is not ``None`` for context-free bandits, the predictions returned will be a
+            list of the same length as contexts.
 
         Returns
         -------
@@ -1196,8 +1175,7 @@ def predict(self,
         ------
         TypeError:  Contexts is not given as ``None``, list, numpy array, pandas series or data frames.
 
-        ValueError: Predicting with contexts data when there is no contextual policy.
-        ValueError: Contextual policy when predicting with no contexts data.
+        ValueError: Prediction with context policy requires context data.
         """
 
         # Check that fit is called before
@@ -1224,7 +1202,8 @@ def predict_expectations(self,
         ----------
         contexts : Union[None, List[Num], List[List[Num]], np.ndarray, pd.Series, pd.DataFrame]
             The context for the expected rewards. Default value is None.
-            Contexts should be ``None`` for context-free bandits and is required for contextual bandits.
+            If contexts is not ``None`` for context-free bandits, the predicted expectations returned will be a
+            list of the same length as contexts.
 
         Returns
         -------
@@ -1234,8 +1213,7 @@ def predict_expectations(self,
         ------
         TypeError:  Contexts is not given as ``None``, list, numpy array or pandas data frames.
 
-        ValueError: Predicting with contexts data when there is no contextual policy.
-        ValueError: Contextual policy when predicting with no contexts data.
+        ValueError: Prediction with context policy requires context data.
         """
 
         # Check that fit is called before
@@ -1382,7 +1360,8 @@ def _validate_predict_args(self, contexts):
             check_true(contexts is not None, ValueError("Prediction with context policy requires context data."))
             MAB._validate_context_type(contexts)
         else:
-            check_true(contexts is None, ValueError("Prediction with no context policy cannot handle context data."))
+            if contexts is not None:
+                MAB._validate_context_type(contexts)
 
     @staticmethod
     def _validate_context_type(contexts):
@@ -1486,7 +1465,9 @@ def __convert_context(self, contexts, decisions=None) -> Union[None, np.ndarray]
             else:  # For predictions, compare the shape to the stored context history
 
                 # We need to find out the number of features (to distinguish Series shape)
-                if isinstance(self.learning_policy, (LearningPolicy.LinGreedy, LearningPolicy.LinTS, LearningPolicy.LinUCB)):
+                if isinstance(self.learning_policy, (LearningPolicy.LinGreedy,
+                                                     LearningPolicy.LinTS,
+                                                     LearningPolicy.LinUCB)):
                     first_arm = self.arms[0]
                     if isinstance(self._imp, _Linear):
                         num_features = self._imp.arm_to_model[first_arm].beta.size
diff --git a/mabwiser/neighbors.py b/mabwiser/neighbors.py
index 7e3ce75..1fd7b09 100755
--- a/mabwiser/neighbors.py
+++ b/mabwiser/neighbors.py
@@ -64,12 +64,12 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.n
         self.contexts = np.concatenate((self.contexts, contexts))
         self.rewards = np.concatenate((self.rewards, rewards))
 
-    def predict(self, contexts: np.ndarray = None):
+    def predict(self, contexts: np.ndarray = None) -> Union[Arm, List[Arm]]:
 
         # Return predict within the neighborhood
         return self._parallel_predict(contexts, is_predict=True)
 
-    def predict_expectations(self, contexts: np.ndarray = None):
+    def predict_expectations(self, contexts: np.ndarray = None) -> Union[Dict[Arm, Num], List[Dict[Arm, Num]]]:
 
         # Return predict expectations within the neighborhood
         return self._parallel_predict(contexts, is_predict=False)
diff --git a/mabwiser/popularity.py b/mabwiser/popularity.py
index 99f0fcf..007829b 100644
--- a/mabwiser/popularity.py
+++ b/mabwiser/popularity.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Dict, List, Optional, NoReturn
+from typing import Dict, List, Optional, NoReturn, Union
 import numpy as np
 
 from mabwiser.greedy import _EpsilonGreedy
@@ -31,20 +31,31 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.n
         # Make sure expectations sum up to 1 like probabilities
         self._normalize_expectations()
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
 
-        # Return the first arm with maximum expectation
-        return argmax(self.predict_expectations())
+        # Return the arm with maximum expectation
+        expectations = self.predict_expectations(contexts)
+        if isinstance(expectations, dict):
+            return argmax(expectations)
+        else:
+            return [argmax(exp) for exp in expectations]
 
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
 
         # Return a random value between 0 and 1 for each arm that is "proportional" to the
         # expectation of the arm and sums to 1 by sampling from a Dirichlet distribution.
         # The Dirichlet distribution can be seen as a multivariate generalization of the Beta distribution.
         # Add a very small epsilon to ensure each of the expectations is positive.
         # TODO: this would not work for negative rewards!
-        expectations = [v + np.finfo(float).eps for v in self.arm_to_expectation.values()]
-        return dict(zip(self.arm_to_expectation.keys(), self.rng.dirichlet(expectations))).copy()
+        size = 1 if contexts is None else len(contexts)
+        alpha = [v + np.finfo(float).eps for v in self.arm_to_expectation.values()]
+        dirichlet_random_values = self.rng.dirichlet(alpha, size)
+        expectations = [dict(zip(self.arm_to_expectation.keys(), exp)).copy() for exp in dirichlet_random_values]
+        if size == 1:
+            return expectations[0]
+        else:
+            return expectations
 
     def _normalize_expectations(self):
         # TODO: this would not work for negative rewards!
diff --git a/mabwiser/rand.py b/mabwiser/rand.py
index bbe246d..562c485 100755
--- a/mabwiser/rand.py
+++ b/mabwiser/rand.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Callable, Dict, List, NoReturn, Optional
+from typing import Callable, Dict, List, NoReturn, Optional, Union
 
 import numpy as np
 
@@ -20,15 +20,28 @@ def fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.ndarray =
     def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.ndarray = None) -> NoReturn:
         pass
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
-
-        # Return the first arm with maximum expectation
-        return argmax(self.predict_expectations())
-
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
-
-        # Return a random expectation (between 0 and 1) for each arm
-        return dict((arm, self.rng.rand()) for arm in self.arms).copy()
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
+
+        # Return the arm with maximum expectation
+        expectations = self.predict_expectations(contexts)
+        if isinstance(expectations, dict):
+            return argmax(expectations)
+        else:
+            return [argmax(exp) for exp in expectations]
+
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
+
+        # Return a random expectation (between 0 and 1) for each arm.
+        # If contexts is None or has length of 1 generate single arm to expectations,
+        # otherwise use vectorized functions to generate a list of arm to expectations with same length as contexts.
+        size = 1 if contexts is None else len(contexts)
+        random_values = self.rng.rand((size, len(self.arms)))
+        expectations = [dict(zip(self.arms, exp)).copy() for exp in random_values]
+        if size == 1:
+            return expectations[0]
+        else:
+            return expectations
 
     def warm_start(self, arm_to_features: Dict[Arm, List[Num]], distance_quantile: float):
         pass
diff --git a/mabwiser/softmax.py b/mabwiser/softmax.py
index d992e80..d74ce5d 100755
--- a/mabwiser/softmax.py
+++ b/mabwiser/softmax.py
@@ -42,19 +42,30 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray,
         self._parallel_fit(decisions, rewards)
         self._expectation_operation()
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
 
-        # Return the first arm with maximum expectation
-        return argmax(self.predict_expectations())
+        # Return the arm with maximum expectation
+        expectations = self.predict_expectations(contexts)
+        if isinstance(expectations, dict):
+            return argmax(expectations)
+        else:
+            return [argmax(exp) for exp in expectations]
 
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
 
         # Return a random value between 0 and 1 for each arm that is "proportional" to the
         # expectation of the arm and sums to 1 by sampling from a Dirichlet distribution.
         # The Dirichlet distribution can be seen as a multivariate generalization of the Beta distribution.
         # Add a very small epsilon to ensure each of the expectations is positive.
-        expectations = [v + np.finfo(float).eps for v in self.arm_to_expectation.values()]
-        return dict(zip(self.arm_to_expectation.keys(), self.rng.dirichlet(expectations))).copy()
+        size = 1 if contexts is None else len(contexts)
+        alpha = [v + np.finfo(float).eps for v in self.arm_to_expectation.values()]
+        dirichlet_random_values = self.rng.dirichlet(alpha, size)
+        expectations = [dict(zip(self.arm_to_expectation.keys(), exp)).copy() for exp in dirichlet_random_values]
+        if size == 1:
+            return expectations[0]
+        else:
+            return expectations
 
     def _copy_arms(self, cold_arm_to_warm_arm):
         for cold_arm, warm_arm in cold_arm_to_warm_arm.items():
diff --git a/mabwiser/thompson.py b/mabwiser/thompson.py
index 84cbe5d..48454c2 100755
--- a/mabwiser/thompson.py
+++ b/mabwiser/thompson.py
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
-from typing import Dict, List, NoReturn, Optional, Callable
+from typing import Callable, Dict, List, NoReturn, Optional, Union
 
 import numpy as np
 
@@ -45,20 +45,31 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray,
         # Calculate fit
         self._parallel_fit(decisions, rewards)
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
 
-        # Return the arm with maximum expectation. If multiple max value exists, return the first one
-        return argmax(self.predict_expectations())
+        # Return the arm with maximum expectation
+        expectations = self.predict_expectations(contexts)
+        if isinstance(expectations, dict):
+            return argmax(expectations)
+        else:
+            return [argmax(exp) for exp in expectations]
 
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
 
-        # Expectation of each arm is a random sample from beta distribution with  success and fail counters
+        # Expectation of each arm is a random sample from beta distribution with success and fail counters.
+        # If contexts is None or has length of 1 generate single arm to expectations,
+        # otherwise use vectorized functions to generate a list of arm to expectations with same length as contexts.
+        size = 1 if contexts is None else len(contexts)
+        arm_to_random_beta = dict()
         for arm in self.arm_to_expectation:
-            self.arm_to_expectation[arm] = self.rng.beta(self.arm_to_success_count[arm],
-                                                         self.arm_to_fail_count[arm])
-
-        # Return a copy of expectations dictionary from arms (key) to expectations (values)
-        return self.arm_to_expectation.copy()
+            arm_to_random_beta[arm] = self.rng.beta(self.arm_to_success_count[arm], self.arm_to_fail_count[arm], size)
+        arm_to_expectation = [{arm: arm_to_random_beta[arm][i] for arm in self.arms} for i in range(size)]
+        self.arm_to_expectation = arm_to_expectation[-1]
+        if size == 1:
+            return arm_to_expectation[0]
+        else:
+            return arm_to_expectation
 
     def _copy_arms(self, cold_arm_to_warm_arm):
         for cold_arm, warm_arm in cold_arm_to_warm_arm.items():
diff --git a/mabwiser/treebandit.py b/mabwiser/treebandit.py
index e7e59a4..09bd333 100644
--- a/mabwiser/treebandit.py
+++ b/mabwiser/treebandit.py
@@ -59,11 +59,11 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray, contexts: np.n
         # Calculate fit
         self._parallel_fit(decisions, rewards, contexts)
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
+    def predict(self, contexts: np.ndarray = None) -> Union[Arm, List[Arm]]:
 
         return self._parallel_predict(contexts, is_predict=True)
 
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: np.ndarray = None) -> Union[Dict[Arm, Num], List[Dict[Arm, Num]]]:
 
         return self._parallel_predict(contexts, is_predict=False)
 
diff --git a/mabwiser/ucb.py b/mabwiser/ucb.py
index e15c461..90011e1 100755
--- a/mabwiser/ucb.py
+++ b/mabwiser/ucb.py
@@ -3,7 +3,7 @@
 
 import math
 from copy import deepcopy
-from typing import Callable, Dict, List, NoReturn, Optional
+from typing import Callable, Dict, List, NoReturn, Optional, Union
 
 import numpy as np
 
@@ -46,15 +46,23 @@ def partial_fit(self, decisions: np.ndarray, rewards: np.ndarray,
         # Calculate fit
         self._parallel_fit(decisions, rewards)
 
-    def predict(self, contexts: np.ndarray = None) -> Arm:
+    def predict(self, contexts: Optional[np.ndarray] = None) -> Union[Arm, List[Arm]]:
 
-        # Return the first arm with maximum expectation
-        return argmax(self.predict_expectations())
+        # Return the arm with maximum expectation
+        expectations = self.predict_expectations(contexts)
+        if isinstance(expectations, dict):
+            return argmax(expectations)
+        else:
+            return [argmax(exp) for exp in expectations]
 
-    def predict_expectations(self, contexts: np.ndarray = None) -> Dict[Arm, Num]:
+    def predict_expectations(self, contexts: Optional[np.ndarray] = None) -> Union[Dict[Arm, Num],
+                                                                                   List[Dict[Arm, Num]]]:
 
         # Return a copy of expectations dictionary from arms (key) to expectations (values)
-        return self.arm_to_expectation.copy()
+        if contexts is None or len(contexts) == 1:
+            return self.arm_to_expectation.copy()
+        else:
+            return [self.arm_to_expectation.copy() for _ in range(len(contexts))]
 
     def _copy_arms(self, cold_arm_to_warm_arm):
         for cold_arm, warm_arm in cold_arm_to_warm_arm.items():
diff --git a/mabwiser/utils.py b/mabwiser/utils.py
index 35d7879..77996fa 100755
--- a/mabwiser/utils.py
+++ b/mabwiser/utils.py
@@ -77,16 +77,24 @@ def __init__(self, seed: int):
         self.rng = None
 
     @abc.abstractmethod
-    def rand(self):
-        """ Return a single float in the range [0, 1)
+    def rand(self, size=None):
+        """ Return return values in range [0, 1) with a given shape.
+
+            Parameters
+            ----------
+            size : int or tuple of ints, optional
+                Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
+                ``m * n * k`` samples are drawn.  Default is None, in which case a
+                single value is returned.
+
             Returns
             -------
-            out : Return a single float in the range [0, 1)
+            out : Array of random floats of shape size (unless size=None, in which case a single float is returned).
         """
         pass
 
     @abc.abstractmethod
-    def randint(self, low: int, high: int = None, size: int = None):
+    def randint(self, low: int, high: int = None, size=None):
         """ Return random integers from low (inclusive) to high (exclusive).
             Return random integers from the “discrete uniform” distribution
             in the “half-open” interval [low, high).
@@ -141,7 +149,7 @@ def choice(self, a: Union[int, Iterable[int]], size: Union[int, Tuple[int]] = No
         pass
 
     @abc.abstractmethod
-    def beta(self, alpha: int, beta: int):
+    def beta(self, alpha: int, beta: int, size=None):
         """ Return a sample from a Beta distribution.
 
             Parameters
@@ -150,16 +158,21 @@ def beta(self, alpha: int, beta: int):
                     Alpha, positive (>0).
             beta : float or array_like of floats
                     Beta, positive (>0).
+            size : int or tuple of ints, optional
+                    Output shape.
+                    If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn.
+                    If size is None (default), a single value is returned if a and b are both scalars.
+                    Otherwise, np.broadcast(a, b).size samples are drawn.
 
             Returns
             -------
-            out : scalar
-                Drawn sample from the parameterized beta distribution.
+            out : scalar or ndarray
+               Drawn samples from the parameterized beta distribution.
         """
         pass
 
     @abc.abstractmethod
-    def standard_normal(self, size):
+    def standard_normal(self, size=None):
         """ Draw samples from a standard Normal distribution (mean=0, stdev=1).
 
             Parameters
@@ -225,8 +238,8 @@ def __init__(self, seed):
         super().__init__(seed)
         self.rng = np.random.default_rng(self.seed)
 
-    def rand(self):
-        return self.rng.random()
+    def rand(self, size=None):
+        return self.rng.random(size)
 
     def randint(self, low: int, high: int = None, size: int = None):
         return self.rng.integers(low=low, high=high, size=size)
@@ -234,10 +247,10 @@ def randint(self, low: int, high: int = None, size: int = None):
     def choice(self, a: Union[int, Iterable[int]], size: Union[int, Tuple[int]] = None, p: Iterable[float] = None):
         return self.rng.choice(a=a, size=size, p=p)
 
-    def beta(self, num_success: int, num_failure: int):
-        return self.rng.beta(num_success, num_failure)
+    def beta(self, num_success: int, num_failure: int, size=None):
+        return self.rng.beta(num_success, num_failure, size)
 
-    def standard_normal(self, size):
+    def standard_normal(self, size=None):
         return self.rng.standard_normal(size)
 
     def multivariate_normal(self, mean: Union[np.ndarray, List[float]],
diff --git a/tests/test_base.py b/tests/test_base.py
index 55eaf66..c50175f 100755
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -114,7 +114,7 @@ def predict(arms: List[Arm],
         else:
 
             # Return: expectations(s) and the MAB instance
-            expectations = [mab.predict_expectations(contexts)for _ in range(num_run)]
+            expectations = [mab.predict_expectations(contexts) for _ in range(num_run)]
             return expectations[0] if num_run == 1 else expectations, mab
 
     @staticmethod
diff --git a/tests/test_greedy.py b/tests/test_greedy.py
index ad6061e..5272853 100755
--- a/tests/test_greedy.py
+++ b/tests/test_greedy.py
@@ -383,3 +383,24 @@ def test_warm_start(self):
         # Warm start
         mab.warm_start(arm_to_features={1: [0, 1], 2: [0, 0], 3: [0.5, 0.5]}, distance_quantile=0.5)
         self.assertDictEqual(mab._imp.arm_to_expectation, {1: 0.5, 2: 0.0, 3: 0.5})
+
+    def test_greedy_contexts(self):
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.EpsilonGreedy(epsilon=0.5),
+                                 contexts=[[]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [1, 1, 1, 1, 1, 1, 3, 1, 1, 1])
+
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.EpsilonGreedy(epsilon=0.5),
+                                 contexts=[[1, 2, 3]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [1, 1, 1, 1, 1, 1, 3, 1, 1, 1])
diff --git a/tests/test_invalid.py b/tests/test_invalid.py
index 1ffc811..e0564f4 100755
--- a/tests/test_invalid.py
+++ b/tests/test_invalid.py
@@ -623,13 +623,9 @@ def test_invalid_jobs(self):
                          is_predict=True,
                          n_jobs=0)
 
-    def test_invalid_add_arm_scaler(self):
-
-        scaler = StandardScaler()
-        arm_to_scaler = {0: deepcopy(scaler), 1: deepcopy(scaler)}
-        mab = MAB([0, 1], LearningPolicy.LinUCB(arm_to_scaler=arm_to_scaler))
+    def test_invalid_scale(self):
         with self.assertRaises(TypeError):
-            mab.add_arm(2, scaler=deepcopy(scaler))
+            mab = MAB([0, 1], LearningPolicy.LinUCB(scale=1))
 
     def test_convert_array_invalid(self):
         df = pd.DataFrame({'a': [1, 1, 1, 1, 1]})
diff --git a/tests/test_lingreedy.py b/tests/test_lingreedy.py
index 9771bcd..bb1ae05 100644
--- a/tests/test_lingreedy.py
+++ b/tests/test_lingreedy.py
@@ -514,6 +514,65 @@ def test_fit_twice(self):
         b_4 = mab._imp.arm_to_model[4].beta
         self.assertEqual(b_4[0], 0)
 
+    def test_fit_twice_scale(self):
+
+        arm, mab = self.predict(arms=[1, 2, 3, 4],
+                                decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                learning_policy=LearningPolicy.LinGreedy(scale=True),
+                                context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                 [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                 [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                 [0, 2, 1, 0, 0]],
+                                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                seed=123456,
+                                num_run=1,
+                                is_predict=True)
+        self.assertEqual(arm, [3, 1])
+        self.assertAlmostEqual(mab._imp.arm_to_model[1].beta[0], -0.1520794283674759)
+        self.assertAlmostEqual(mab._imp.arm_to_model[2].beta[0], 0)
+        self.assertAlmostEqual(mab._imp.arm_to_model[3].beta[0], -0.008110550702115856)
+        self.assertAlmostEqual(mab._imp.arm_to_model[4].beta[0], 0)
+
+        # Fit again
+        mab.fit(decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                          [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                          [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                          [0, 2, 1, 0, 0]])
+        self.assertEqual(arm, [3, 1])
+        self.assertAlmostEqual(mab._imp.arm_to_model[1].beta[0], -0.1520794283674759)
+        self.assertAlmostEqual(mab._imp.arm_to_model[2].beta[0], 0)
+        self.assertAlmostEqual(mab._imp.arm_to_model[3].beta[0], -0.008110550702115856)
+        self.assertAlmostEqual(mab._imp.arm_to_model[4].beta[0], 0)
+
+    def test_add_arm_scale(self):
+        arm, mab = self.predict(arms=[1, 2, 3],
+                                decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                learning_policy=LearningPolicy.LinGreedy(scale=True),
+                                context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                 [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                 [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                 [0, 2, 1, 0, 0]],
+                                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                seed=123456,
+                                num_run=1,
+                                is_predict=True)
+        mab.add_arm(4)
+        mab.partial_fit(decisions=[1, 1, 4, 4],
+                        rewards=[0, 1, 1, 1],
+                        contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0], [0, 2, 2, 3, 5]])
+        self.assertEqual(mab._imp.arm_to_model[1].scaler.n_samples_seen_, 5)
+        self.assertEqual(mab._imp.arm_to_model[2].scaler.n_samples_seen_, 2)
+        self.assertEqual(mab._imp.arm_to_model[3].scaler.n_samples_seen_, 5)
+        self.assertEqual(mab._imp.arm_to_model[4].scaler.n_samples_seen_, 2)
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[1].scaler.mean_), [0.4, 0.8, 1.4, 1.6, 2.4])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[2].scaler.mean_), [0.5, 2.5, 1.5, 2.0, 3.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[3].scaler.mean_), [0.2, 1.2, 1.6, 1.6, 2.6])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[4].scaler.mean_), [0.0, 1.0, 1.5, 1.5, 2.5])
+
     def test_partial_fit(self):
 
         arm, mab = self.predict(arms=[1, 2, 3, 4],
@@ -592,7 +651,86 @@ def test_alpha0_nearest5(self):
         self.assertEqual(len(arm), 3)
         self.assertEqual(arm, [[3, 2], [3, 3], [3, 3]])
 
-    def test_scaler(self):
+    def test_scaler_fit(self):
+        exp, mab = self.predict(arms=[1, 2, 3],
+                                decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                learning_policy=LearningPolicy.LinGreedy(scale=True),
+                                context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                 [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                 [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                 [0, 2, 1, 0, 0]],
+                                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                seed=123456,
+                                num_run=1,
+                                is_predict=False)
+        self.assertEqual(mab._imp.arm_to_model[1].scaler.n_samples_seen_, 3)
+        self.assertEqual(mab._imp.arm_to_model[2].scaler.n_samples_seen_, 2)
+        self.assertEqual(mab._imp.arm_to_model[3].scaler.n_samples_seen_, 5)
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[1].scaler.mean_), [1/3, 2/3,  4/3, 4/3, 2.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[2].scaler.mean_), [0.5, 2.5, 1.5, 2.0, 3.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[3].scaler.mean_), [0.2, 1.2, 1.6, 1.6, 2.6])
+
+    def test_scaler_fit_twice(self):
+        exp, mab = self.predict(arms=[1, 2, 3],
+                                decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                learning_policy=LearningPolicy.LinGreedy(scale=True),
+                                context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                 [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                 [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                 [0, 2, 1, 0, 0]],
+                                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                seed=123456,
+                                num_run=1,
+                                is_predict=False)
+        self.assertEqual(mab._imp.arm_to_model[1].scaler.n_samples_seen_, 3)
+        self.assertEqual(mab._imp.arm_to_model[2].scaler.n_samples_seen_, 2)
+        self.assertEqual(mab._imp.arm_to_model[3].scaler.n_samples_seen_, 5)
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[1].scaler.mean_), [1/3, 2/3,  4/3, 4/3, 2.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[2].scaler.mean_), [0.5, 2.5, 1.5, 2.0, 3.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[3].scaler.mean_), [0.2, 1.2, 1.6, 1.6, 2.6])
+
+        mab.fit(decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                          [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                          [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                          [0, 2, 1, 0, 0]])
+        self.assertEqual(mab._imp.arm_to_model[1].scaler.n_samples_seen_, 3)
+        self.assertEqual(mab._imp.arm_to_model[2].scaler.n_samples_seen_, 2)
+        self.assertEqual(mab._imp.arm_to_model[3].scaler.n_samples_seen_, 5)
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[1].scaler.mean_), [1/3, 2/3,  4/3, 4/3, 2.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[2].scaler.mean_), [0.5, 2.5, 1.5, 2.0, 3.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[3].scaler.mean_), [0.2, 1.2, 1.6, 1.6, 2.6])
+
+    def test_scaler_partial_fit(self):
+        exp, mab = self.predict(arms=[1, 2, 3],
+                                decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                learning_policy=LearningPolicy.LinGreedy(scale=True),
+                                context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                 [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                 [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                 [0, 2, 1, 0, 0]],
+                                contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                seed=123456,
+                                num_run=1,
+                                is_predict=False)
+        mab.partial_fit(decisions=[3, 3, 3, 2, 2, 1, 1, 1, 1, 1],
+                        rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                        contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                  [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                  [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                  [0, 2, 1, 0, 0]])
+        self.assertEqual(mab._imp.arm_to_model[1].scaler.n_samples_seen_, 8)
+        self.assertEqual(mab._imp.arm_to_model[2].scaler.n_samples_seen_, 4)
+        self.assertEqual(mab._imp.arm_to_model[3].scaler.n_samples_seen_, 8)
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[1].scaler.mean_), [0.25, 1.0, 1.5, 1.5, 2.375])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[2].scaler.mean_), [0.5, 2.5, 1.5, 2.0, 3.0])
+        self.assertListAlmostEqual(list(mab._imp.arm_to_model[3].scaler.mean_), [0.25, 1.0, 1.5, 1.5, 2.375])
+
+    def test_scaler_predictions(self):
 
         arms = [1, 2, 3]
         context_history = np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
@@ -614,7 +752,7 @@ def test_scaler(self):
         exp, mab = self.predict(arms=arms,
                                 decisions=decisions,
                                 rewards=rewards,
-                                learning_policy=LearningPolicy.LinGreedy(arm_to_scaler=arm_to_scaler),
+                                learning_policy=LearningPolicy.LinGreedy(scale=True),
                                 context_history=context_history,
                                 contexts=contexts,
                                 seed=123456,
@@ -641,6 +779,23 @@ def test_scaler(self):
             for i in range(len(contexts)):
                 self.assertEqual(exp[i][arm], exp_check[i][arm])
 
+    def test_unused_arm_scale(self):
+
+        arms, mab = self.predict(arms=[1, 2, 3, 4],
+                                 decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                 rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.LinGreedy(l2_lambda=1, scale=True),
+                                 context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                  [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                  [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                  [0, 2, 1, 0, 0]],
+                                 contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+
+        self.assertEqual(arms, [3, 1])
+
     def test_add_arm(self):
         arm, mab = self.predict(arms=[1, 2, 3],
                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3, 1],
diff --git a/tests/test_lints.py b/tests/test_lints.py
index 7200c88..f76e18b 100644
--- a/tests/test_lints.py
+++ b/tests/test_lints.py
@@ -7,100 +7,79 @@
 import pandas as pd
 from sklearn.preprocessing import StandardScaler
 
-from mabwiser.mab import LearningPolicy, NeighborhoodPolicy
+from mabwiser.mab import LearningPolicy
 from tests.test_base import BaseTest
 
 
 class LinTSTest(BaseTest):
 
     def test_alpha0_0001(self):
-        scaler = StandardScaler()
-        context_history = np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
-                                                 [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
-                                                 [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
-                                                 [0, 2, 1, 0, 0]])
-        scaler.fit(context_history)
-
         arm, mab = self.predict(arms=[1, 2, 3],
                                 decisions=[1, 1, 1, 2, 2, 2, 3, 3, 3, 1],
                                 rewards=[0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
-                                learning_policy=LearningPolicy.LinTS(alpha=0.0001,
-                                                                     arm_to_scaler={1: scaler, 2: scaler, 3: scaler}),
-                                context_history=context_history,
+                                learning_policy=LearningPolicy.LinTS(alpha=0.0001, scale=True),
+                                context_history=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                         [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                         [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                         [0, 2, 1, 0, 0]]),
                                 contexts=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]]),
                                 seed=123456,
                                 num_run=3,
                                 is_predict=True)
 
         self.assertEqual(len(arm), 3)
-        self.assertEqual(arm, [[3, 2], [3, 2], [3, 2]])
+        self.assertEqual(arm, [[2, 3], [2, 3], [3, 3]])
 
     def test_alpha0_0001_expectations(self):
-        scaler = StandardScaler()
-        context_history = np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
-                                    [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
-                                    [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
-                                    [0, 2, 1, 0, 0]])
-        scaler.fit(context_history)
-
         exps, mab = self.predict(arms=[1, 2, 3],
                                  decisions=[1, 1, 1, 2, 2, 2, 3, 3, 3, 1],
                                  rewards=[0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
-                                 learning_policy=LearningPolicy.LinTS(alpha=0.0001,
-                                                                      arm_to_scaler={1: scaler, 2: scaler, 3: scaler}),
-                                 context_history=context_history,
+                                 learning_policy=LearningPolicy.LinTS(alpha=0.0001, scale=True),
+                                 context_history=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                          [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                          [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                          [0, 2, 1, 0, 0]]),
                                  contexts=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]]),
                                  seed=123456,
                                  num_run=1,
                                  is_predict=False)
 
-        self.assertListAlmostEqual(exps[0].values(), [-0.2852569197274552, 0.00016338898274687822, 0.5735155717025586])
-        self.assertListAlmostEqual(exps[1].values(), [-0.1487907884038659, -7.802191785859073e-05,
-                                                      -0.0008484936824155465])
+        self.assertListAlmostEqual(exps[0].values(),
+                                   [-0.23459369004297587, 0.0002702455674444537, 4.588547880979047e-05])
+        self.assertListAlmostEqual(exps[1].values(),
+                                   [-0.192811601170233, -2.3415795345448245e-05, 0.00016619626256880228])
 
     def test_alpha1(self):
-        scaler = StandardScaler()
-        context_history = np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
-                                    [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
-                                    [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
-                                    [0, 2, 1, 0, 0]])
-        scaler.fit(context_history)
-
         arm, mab = self.predict(arms=[1, 2, 3],
                                 decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
                                 rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
-                                learning_policy=LearningPolicy.LinTS(alpha=1,
-                                                                     arm_to_scaler={1: scaler, 2: scaler, 3: scaler}),
-                                context_history=context_history,
+                                learning_policy=LearningPolicy.LinTS(alpha=1),
+                                context_history=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                         [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                         [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                         [0, 2, 1, 0, 0]]),
                                 contexts=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]]),
                                 seed=123456,
                                 num_run=1,
                                 is_predict=True)
-
         self.assertEqual(len(arm), 2)
         self.assertEqual(arm, [2, 3])
 
     def test_alpha1_expectations(self):
-        scaler = StandardScaler()
-        context_history = np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
-                                    [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
-                                    [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
-                                    [0, 2, 1, 0, 0]])
-        scaler.fit(context_history)
-
         exps, mab = self.predict(arms=[1, 2, 3],
                                  decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
                                  rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
-                                 learning_policy=LearningPolicy.LinTS(alpha=1,
-                                                                      arm_to_scaler={1: scaler, 2: scaler, 3: scaler}),
-                                 context_history=context_history,
+                                 learning_policy=LearningPolicy.LinTS(alpha=1),
+                                 context_history=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                           [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                           [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                           [0, 2, 1, 0, 0]]),
                                  contexts=np.array([[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]]),
                                  seed=123456,
                                  num_run=1,
                                  is_predict=False)
-
-        self.assertListAlmostEqual(exps[0].values(), [-0.6846042491588905, 1.8728586982060706, 0.39597711947956443])
-        self.assertListAlmostEqual(exps[1].values(), [-0.9156881567627143, -1.01000793116177, 1.6774048483779203])
+        self.assertListAlmostEqual(exps[0].values(), [-0.6029872358950072, 3.105765259323796, 1.0208598325762464])
+        self.assertListAlmostEqual(exps[1].values(), [0.572141413757231, 0.45473267178654997, 1.773376616755168])
 
     def test_np(self):
 
@@ -350,6 +329,23 @@ def test_lints_t10(self):
         self.assertEqual(len(arm), 4)
         self.assertEqual(arm, [[b, b], [b, b], [b, b], [b, b]])
 
+    def test_unused_arm_scale(self):
+
+        arms, mab = self.predict(arms=[1, 2, 3, 4],
+                                 decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                 rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.LinTS(alpha=1, scale=True),
+                                 context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                  [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                  [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                  [0, 2, 1, 0, 0]],
+                                 contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+
+        self.assertEqual(arms, [4, 3])
+
     def test_unused_arm(self):
 
         exps, mab = self.predict(arms=[1, 2, 3, 4],
@@ -385,7 +381,6 @@ def test_unused_arm2(self):
                                  num_run=1,
                                  is_predict=True)
 
-
         self.assertEqual(arms, [4, 3])
 
     def test_unused_arm_scaled(self):
diff --git a/tests/test_linucb.py b/tests/test_linucb.py
index 697363a..00930b6 100755
--- a/tests/test_linucb.py
+++ b/tests/test_linucb.py
@@ -8,7 +8,7 @@
 import pandas as pd
 from sklearn.preprocessing import StandardScaler
 
-from mabwiser.mab import MAB, LearningPolicy, NeighborhoodPolicy
+from mabwiser.mab import LearningPolicy, NeighborhoodPolicy
 from tests.test_base import BaseTest
 
 
@@ -559,7 +559,7 @@ def test_scaler(self):
         exp, mab = self.predict(arms=arms,
                                 decisions=decisions,
                                 rewards=rewards,
-                                learning_policy=LearningPolicy.LinUCB(arm_to_scaler=arm_to_scaler),
+                                learning_policy=LearningPolicy.LinUCB(scale=True),
                                 context_history=context_history,
                                 contexts=contexts,
                                 seed=123456,
@@ -586,13 +586,22 @@ def test_scaler(self):
             for i in range(len(contexts)):
                 self.assertAlmostEqual(exp[i][arm], exp_check[i][arm])
 
-    @staticmethod
-    def test_add_arm_scaler():
-        scaler = StandardScaler()
-        scaler.fit(np.array([[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]]).astype('float64'))
-        arm_to_scaler = {0: deepcopy(scaler), 1: deepcopy(scaler)}
-        mab = MAB([0, 1], LearningPolicy.LinUCB(arm_to_scaler=arm_to_scaler))
-        mab.add_arm(2, scaler=deepcopy(scaler))
+    def test_unused_arm_scale(self):
+
+        arms, mab = self.predict(arms=[1, 2, 3, 4],
+                                 decisions=[1, 1, 1, 2, 2, 3, 3, 3, 3, 3],
+                                 rewards=[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.LinUCB(scale=True),
+                                 context_history=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1], [0, 0, 1, 0, 0],
+                                                  [0, 2, 2, 3, 5], [1, 3, 1, 1, 1], [0, 0, 0, 0, 0],
+                                                  [0, 1, 4, 3, 5], [0, 1, 2, 4, 5], [1, 2, 1, 1, 3],
+                                                  [0, 2, 1, 0, 0]],
+                                 contexts=[[0, 1, 2, 3, 5], [1, 1, 1, 1, 1]],
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+
+        self.assertEqual(arms, [4, 2])
 
     def test_add_arm(self):
         arm, mab = self.predict(arms=[1, 2, 3],
diff --git a/tests/test_mab.py b/tests/test_mab.py
index 907cf9a..f330493 100755
--- a/tests/test_mab.py
+++ b/tests/test_mab.py
@@ -2,11 +2,9 @@
 
 import os
 import pickle
-from copy import deepcopy
 
 import numpy as np
 import pandas as pd
-from sklearn.preprocessing import StandardScaler
 
 from mabwiser.mab import MAB, LearningPolicy, NeighborhoodPolicy
 from tests.test_base import BaseTest
@@ -46,17 +44,11 @@ def test_learning_policy_values(self):
         mab = MAB([0, 1], lp)
         self.assertEqual(lp.epsilon, mab.learning_policy.epsilon)
 
-        data = np.array([[1, 2, 3], [3, 2, 1]])
-        sc = StandardScaler()
-        sc.fit(data)
-        arm_to_scaler = {0: sc, 1: sc}
-
-        lp = LearningPolicy.LinUCB(alpha=2.0, l2_lambda=0.3, arm_to_scaler=arm_to_scaler)
+        lp = LearningPolicy.LinUCB(alpha=2.0, l2_lambda=0.3, scale=True)
         mab = MAB([0, 1], lp)
         self.assertEqual(lp.alpha, mab.learning_policy.alpha)
         self.assertEqual(lp.l2_lambda, mab.learning_policy.l2_lambda)
-        self.assertIs(sc, mab.learning_policy.arm_to_scaler[0])
-        self.assertIs(sc, mab.learning_policy.arm_to_scaler[1])
+        self.assertEqual(lp.scale, mab.learning_policy.scale)
 
         lp = LearningPolicy.Softmax(tau=0.5)
         mab = MAB([0, 1], lp)
diff --git a/tests/test_popularity.py b/tests/test_popularity.py
index 5335804..29d9c7e 100644
--- a/tests/test_popularity.py
+++ b/tests/test_popularity.py
@@ -298,3 +298,24 @@ def test_warm_start(self):
         # Warm start
         mab.warm_start(arm_to_features={1: [0, 1], 2: [0, 0], 3: [0, 1]}, distance_quantile=0.5)
         self.assertDictEqual(mab._imp.arm_to_expectation, {1: 1.0, 2: 0.0, 3: 1.0})
+
+    def test_popularity_contexts(self):
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.Popularity(),
+                                 contexts=[[]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [3, 2, 3, 3, 3, 2, 2, 3, 2, 3])
+
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.Popularity(),
+                                 contexts=[[1, 2, 3]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [3, 2, 3, 3, 3, 2, 2, 3, 2, 3])
diff --git a/tests/test_random.py b/tests/test_random.py
index 82e4586..b6e5296 100755
--- a/tests/test_random.py
+++ b/tests/test_random.py
@@ -73,7 +73,7 @@ def test_add_arm(self):
         arms, mab = self.predict(arms=[1, 2, 3],
                                  decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
                                  rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
-                                 learning_policy=LearningPolicy.EpsilonGreedy(epsilon=0.25),
+                                 learning_policy=LearningPolicy.Random(),
                                  seed=123456,
                                  num_run=4,
                                  is_predict=True)
@@ -86,10 +86,31 @@ def test_remove_arm(self):
         arms, mab = self.predict(arms=[1, 2, 3],
                                  decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
                                  rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
-                                 learning_policy=LearningPolicy.EpsilonGreedy(epsilon=0.25),
+                                 learning_policy=LearningPolicy.Random(),
                                  seed=123456,
                                  num_run=4,
                                  is_predict=True)
         mab.remove_arm(3)
         self.assertTrue(3 not in mab.arms)
         self.assertTrue(3 not in mab._imp.arm_to_expectation)
+
+    def test_random_contexts(self):
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.Random(),
+                                 contexts=[[]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [1, 1, 3, 1, 2, 2, 2, 2, 2, 2])
+
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.Random(),
+                                 contexts=[[1, 2, 3]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [1, 1, 3, 1, 2, 2, 2, 2, 2, 2])
diff --git a/tests/test_ridge.py b/tests/test_ridge.py
index 1d33d5d..fc31b71 100755
--- a/tests/test_ridge.py
+++ b/tests/test_ridge.py
@@ -17,7 +17,7 @@ def test_predict_ridge(self):
         rewards = np.array([3, 3, 1])
         rng = np.random.RandomState(seed=7)
 
-        ridge = _RidgeRegression(rng, l2_lambda=1.0, alpha=1.0, scaler = None)
+        ridge = _RidgeRegression(rng, l2_lambda=1.0, alpha=1.0, scale=False)
 
         ridge.init(context.shape[1])
         ridge.fit(context, rewards)
@@ -31,7 +31,7 @@ def test_predict_ridge_scaler(self):
         scaler = StandardScaler()
         scaler.fit(context.astype('float64'))
 
-        ridge = _RidgeRegression(rng, l2_lambda=1.0, alpha=1.0, scaler = scaler)
+        ridge = _RidgeRegression(rng, l2_lambda=1.0, alpha=1.0, scale=True)
 
         ridge.init(context.shape[1])
         ridge.fit(context, rewards)
diff --git a/tests/test_softmax.py b/tests/test_softmax.py
index c3e0e13..634c7b5 100755
--- a/tests/test_softmax.py
+++ b/tests/test_softmax.py
@@ -389,3 +389,24 @@ def test_warm_start(self):
         mab.warm_start(arm_to_features={1: [0, 1], 2: [0, 0], 3: [0, 1]}, distance_quantile=0.5)
         self.assertDictEqual(mab._imp.arm_to_expectation, {1: 0.38365173119055074, 2: 0.23269653761889864,
                                                            3: 0.38365173119055074})
+
+    def test_softmax_contexts(self):
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.Softmax(),
+                                 contexts=[[]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [3, 2, 3, 3, 3, 2, 2, 3, 2, 3])
+
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.Softmax(),
+                                 contexts=[[1, 2, 3]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [3, 2, 3, 3, 3, 2, 2, 3, 2, 3])
diff --git a/tests/test_thompson.py b/tests/test_thompson.py
index 552a5f7..9aa2fd3 100755
--- a/tests/test_thompson.py
+++ b/tests/test_thompson.py
@@ -514,3 +514,24 @@ def test_warm_start(self):
         mab.warm_start(arm_to_features={1: [0, 1], 2: [0, 0], 3: [0, 1]}, distance_quantile=0.5)
         self.assertDictEqual(mab._imp.arm_to_fail_count, {1: 3, 2: 4, 3: 3})
         self.assertDictEqual(mab._imp.arm_to_success_count, {1: 5, 2: 1, 3: 5})
+
+    def test_ts_contexts(self):
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.ThompsonSampling(),
+                                 contexts=[[]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [3, 1, 1, 1, 1, 3, 1, 1, 1, 1])
+
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.ThompsonSampling(),
+                                 contexts=[[1, 2, 3]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [3, 1, 1, 1, 1, 3, 1, 1, 1, 1])
diff --git a/tests/test_ucb.py b/tests/test_ucb.py
index d1526b6..b57470b 100755
--- a/tests/test_ucb.py
+++ b/tests/test_ucb.py
@@ -407,3 +407,24 @@ def test_warm_start(self):
         # Warm start
         mab.warm_start(arm_to_features={1: [0, 1], 2: [0, 0], 3: [0.5, 0.5]}, distance_quantile=0.5)
         self.assertDictEqual(mab._imp.arm_to_mean, {1: 0.5, 2: 0.0, 3: 0.5})
+
+    def test_ucb_contexts(self):
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.UCB1(),
+                                 contexts=[[]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [2, 2, 2, 2, 2, 2, 2, 2, 2, 2])
+
+        arms, mab = self.predict(arms=[1, 2, 3],
+                                 decisions=[1, 1, 1, 3, 2, 2, 3, 1, 3],
+                                 rewards=[0, 1, 1, 0, 1, 0, 1, 1, 1],
+                                 learning_policy=LearningPolicy.UCB1(),
+                                 contexts=[[1, 2, 3]] * 10,
+                                 seed=123456,
+                                 num_run=1,
+                                 is_predict=True)
+        self.assertEqual(arms, [2, 2, 2, 2, 2, 2, 2, 2, 2, 2])