From 1d1c98ab6efdbd2a40081f2c28ddee32dda094e9 Mon Sep 17 00:00:00 2001
From: Alexander Krotov <krotov@iitp.ru>
Date: Sat, 24 Oct 2020 02:03:14 +0300
Subject: [PATCH] wifi: add ThompsonSamplingWifiManager

---
 CHANGES.html                                  |   1 +
 RELEASE_NOTES                                 |   1 +
 src/wifi/doc/source/wifi-design.rst           |  37 ++
 src/wifi/doc/source/wifi-references.rst       |   2 +
 src/wifi/examples/wifi-manager-example.cc     |   4 +-
 src/wifi/helper/wifi-helper.cc                |   7 +
 .../thompson-sampling-wifi-manager.cc         | 422 ++++++++++++++++++
 .../thompson-sampling-wifi-manager.h          | 135 ++++++
 src/wifi/wscript                              |   2 +
 9 files changed, 609 insertions(+), 2 deletions(-)
 create mode 100644 src/wifi/model/rate-control/thompson-sampling-wifi-manager.cc
 create mode 100644 src/wifi/model/rate-control/thompson-sampling-wifi-manager.h

diff --git a/CHANGES.html b/CHANGES.html
index 2a22a7355..2fa7de792 100644
--- a/CHANGES.html
+++ b/CHANGES.html
@@ -57,6 +57,7 @@ us a note on ns-developers mailing list.</p>
 <li>Added <b>FqCobalt</b> queue disc with L4S features and set associative hash.</li>
 <li>Added <b>FqPIE</b> queue disc with <b>L4S</b> mode</li>
 <li>Added the ability to configure the primary 20 MHz channel for 802.11 devices operating on channels of width greater than 20 MHz.</li>
+<li>Added new <b>ThompsonSamplingWifiManager</b> rate control algorithm.</li>
 </ul>
 <h2>Changes to existing API:</h2>
 <ul>
diff --git a/RELEASE_NOTES b/RELEASE_NOTES
index a43ceb5a9..b36278f1a 100644
--- a/RELEASE_NOTES
+++ b/RELEASE_NOTES
@@ -36,6 +36,7 @@ New user-visible features
 - (wifi) Stations perform TXOP recovery if the transmission of a non-initial MPDU in a TXOP fails
 - (wifi) Stations keep track of the TXOP holder and ignore the NAV when they receive an RTS frame from the TXOP holder
 - (wifi) The TxOkHeader and TxErrHeader trace sources of RegularWifiMac have been obsoleted and replaced by trace sources that better capture the result of a transmission (AckedMpdu, NAckedMpdu, DroppedMpdu, MpduResponseTimeout and PsduResponseTimeout)
+- (wifi) Add ThompsonSamplingWifiManager rate control algorithm.
 - (traffic-control) Added FqCobalt queue disc with L4S features and set associative hash.
 - (traffic-control) Added FqPIE queue disc with L4S mode.
 
diff --git a/src/wifi/doc/source/wifi-design.rst b/src/wifi/doc/source/wifi-design.rst
index f89d3c2ac..748c08018 100644
--- a/src/wifi/doc/source/wifi-design.rst
+++ b/src/wifi/doc/source/wifi-design.rst
@@ -948,6 +948,7 @@ Algorithms in literature:
 * ``AarfcdWifiManager`` [maguolo2008aarfcd]_
 * ``ParfWifiManager`` [akella2007parf]_
 * ``AparfWifiManager`` [chevillat2005aparf]_
+* ``ThompsonSamplingWifiManager`` [krotov2020rate]_
 
 ConstantRateWifiManager
 #######################
@@ -1010,6 +1011,42 @@ With this new default value (i.e. 1e-6), a HE STA moving away from a HE AP has
 smooth throughput decrease (whereas with 1e-5, better performance was seen further
 away, which is not "ideal").
 
+ThompsonSamplingWifiManager
+###########################
+
+Thompson Sampling (TS) is a classical solution to the Multi-Armed
+Bandit problem.  `ThompsonSamplingWifiManager` implements a rate
+control algorithm based on TS with the goal of providing a simple
+statistics-based algorithm with a low number of parameters.
+
+The algorithm maintains the number of successful transmissions
+:math:`\alpha_i` and the number of unsuccessful transmissions
+:math:`\beta_i` for each MCS :math:`i`, both of which are initially
+set to zero.
+
+To select MCS for a data frame, the algorithm draws a sample frame
+success rate :math:`\p_i` from the beta distribution with shape
+parameters :math:`(1 + \alpha_i, 1 + \beta_i)` for each MCS and then
+selects MCS with the highest expected throughput calculated as the
+sample frame success rate multiplied by MCS rate.
+
+To account for changing channel conditions, exponential decay is
+applied to :math:`\alpha_i` and :math:`\beta_i`. The rate of
+exponential decay is controlled with the `Decay` attribute which is
+the inverse of the time constant. Default value of 1 Hz results in
+using exponential window with the time constant of 1 second.  Setting
+this value to zero effectively disables exponential decay and can be
+used in static scenarios.
+
+Control frames are always transmitted using the most robust MCS,
+except when the standard specifies otherwise, such as for ACK frames.
+
+As the main goal of this algorithm is to provide a stable baseline, it
+does not take into account backoff overhead, inter-frame spaces and
+aggregation for MCS rate calculation. For an example of a more complex
+statistics-based rate control algorithm used in real devices, consider
+Minstrel-HT described below.
+
 MinstrelWifiManager
 ###################
 
diff --git a/src/wifi/doc/source/wifi-references.rst b/src/wifi/doc/source/wifi-references.rst
index 154a7d37d..40517e4ae 100644
--- a/src/wifi/doc/source/wifi-references.rst
+++ b/src/wifi/doc/source/wifi-references.rst
@@ -58,3 +58,5 @@ References
 .. [erceg2004] \ V. Erceg and L. Schumacher and P. Kyritsi, "Tgn channel models", IEEE 802.11-03/940r4, 2004.
 
 .. [porat2016] \ R. Porat et al., "11ax Evaluation Methodology", IEE P802.11 Wireless LANs, 11-14-0571r3, 2016. 
+
+.. [krotov2020rate] \ A. Krotov, A. Kiryanov, E. Khorov., `Rate Control With Spatial Reuse for Wi-Fi 6 Dense Deployments <https://doi.org/10.1109/ACCESS.2020.3023552>`__, IEEE Access, September 2020
diff --git a/src/wifi/examples/wifi-manager-example.cc b/src/wifi/examples/wifi-manager-example.cc
index 4403e0400..0a65ad7f3 100644
--- a/src/wifi/examples/wifi-manager-example.cc
+++ b/src/wifi/examples/wifi-manager-example.cc
@@ -29,7 +29,7 @@
 //
 // By default, the 802.11a standard using IdealWifiManager is plotted. Several command line
 // arguments can change the following options:
-// --wifiManager (Aarf, Aarfcd, Amrr, Arf, Cara, Ideal, Minstrel, MinstrelHt, Onoe, Rraa)
+// --wifiManager (Aarf, Aarfcd, Amrr, Arf, Cara, Ideal, Minstrel, MinstrelHt, Onoe, Rraa, ThompsonSampling)
 // --standard (802.11a, 802.11b, 802.11g, 802.11n-5GHz, 802.11n-2.4GHz, 802.11ac, 802.11p-10MHz, 802.11p-5MHz)
 // --serverShortGuardInterval and --clientShortGuardInterval (for 802.11n/ac)
 // --serverNss and --clientNss (for 802.11n/ac)
@@ -187,7 +187,7 @@ int main (int argc, char *argv[])
   cmd.AddValue ("serverShortGuardInterval", "Set short guard interval of the server (802.11n/ac/ax) in nanoseconds", serverShortGuardInterval);
   cmd.AddValue ("clientShortGuardInterval", "Set short guard interval of the client (802.11n/ac/ax) in nanoseconds", clientShortGuardInterval);
   cmd.AddValue ("standard", "Set standard (802.11a, 802.11b, 802.11g, 802.11n-5GHz, 802.11n-2.4GHz, 802.11ac, 802.11p-10MHz, 802.11p-5MHz, 802.11ax-5GHz, 802.11ax-2.4GHz)", standard);
-  cmd.AddValue ("wifiManager", "Set wifi rate manager (Aarf, Aarfcd, Amrr, Arf, Cara, Ideal, Minstrel, MinstrelHt, Onoe, Rraa)", wifiManager);
+  cmd.AddValue ("wifiManager", "Set wifi rate manager (Aarf, Aarfcd, Amrr, Arf, Cara, Ideal, Minstrel, MinstrelHt, Onoe, Rraa, ThompsonSampling)", wifiManager);
   cmd.AddValue ("infrastructure", "Use infrastructure instead of adhoc", infrastructure);
   cmd.Parse (argc,argv);
 
diff --git a/src/wifi/helper/wifi-helper.cc b/src/wifi/helper/wifi-helper.cc
index 700f17b6f..9dd223461 100644
--- a/src/wifi/helper/wifi-helper.cc
+++ b/src/wifi/helper/wifi-helper.cc
@@ -23,6 +23,7 @@
 #include "ns3/wifi-net-device.h"
 #include "ns3/minstrel-wifi-manager.h"
 #include "ns3/minstrel-ht-wifi-manager.h"
+#include "ns3/thompson-sampling-wifi-manager.h"
 #include "ns3/ap-wifi-mac.h"
 #include "ns3/ampdu-subframe-header.h"
 #include "ns3/mobility-model.h"
@@ -1034,6 +1035,12 @@ WifiHelper::AssignStreams (NetDeviceContainer c, int64_t stream)
               currentStream += minstrelHt->AssignStreams (currentStream);
             }
 
+          Ptr<ThompsonSamplingWifiManager> thompsonSampling = DynamicCast<ThompsonSamplingWifiManager> (manager);
+          if (thompsonSampling)
+            {
+              currentStream += thompsonSampling->AssignStreams (currentStream);
+            }
+
           //Handle any random numbers in the MAC objects.
           Ptr<WifiMac> mac = wifi->GetMac ();
           Ptr<RegularWifiMac> rmac = DynamicCast<RegularWifiMac> (mac);
diff --git a/src/wifi/model/rate-control/thompson-sampling-wifi-manager.cc b/src/wifi/model/rate-control/thompson-sampling-wifi-manager.cc
new file mode 100644
index 000000000..e4f05df51
--- /dev/null
+++ b/src/wifi/model/rate-control/thompson-sampling-wifi-manager.cc
@@ -0,0 +1,422 @@
+/* -*- Mode:C++; c-file-style:"gnu"; indent-tabs-mode:nil; -*- */
+/*
+ * Copyright (c) 2021 IITP RAS
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License version 2 as
+ * published by the Free Software Foundation;
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+ *
+ * Author: Alexander Krotov <krotov@iitp.ru>
+ */
+
+#include "ns3/log.h"
+#include "ns3/double.h"
+#include "ns3/core-module.h"
+#include "ns3/packet.h"
+
+#include "ns3/wifi-phy.h"
+
+#include "thompson-sampling-wifi-manager.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <fstream>
+#include <iostream>
+#include <string>
+
+namespace ns3 {
+
+/**
+ * A structure containing parameters of a single rate and its
+ * statistics.
+ */
+struct RateStats {
+  WifiMode mode; ///< MCS
+  uint16_t channelWidth; ///< channel width in MHz
+  uint8_t nss; ///< Number of spatial streams
+
+  double success{0.0}; ///< averaged number of successful transmissions
+  double fails{0.0}; ///< averaged number of failed transmissions
+  Time lastDecay{0}; ///< last time exponential decay was applied to this rate
+};
+
+/**
+ * Holds station state and collected statistics.
+ *
+ * This struct extends from WifiRemoteStation to hold additional
+ * information required by ThompsonSamplingWifiManager.
+ */
+struct ThompsonSamplingWifiRemoteStation : public WifiRemoteStation
+{
+  size_t m_nextMode; //!< Mode to select for the next transmission
+  size_t m_lastMode; //!< Most recently used mode, used to write statistics
+
+  std::vector<RateStats> m_mcsStats; //!< Collected statistics
+};
+
+NS_OBJECT_ENSURE_REGISTERED (ThompsonSamplingWifiManager);
+
+NS_LOG_COMPONENT_DEFINE ("ThompsonSamplingWifiManager");
+
+TypeId
+ThompsonSamplingWifiManager::GetTypeId (void)
+{
+  static TypeId tid = TypeId ("ns3::ThompsonSamplingWifiManager")
+    .SetParent<WifiRemoteStationManager> ()
+    .SetGroupName ("Wifi")
+    .AddConstructor<ThompsonSamplingWifiManager> ()
+    .AddAttribute ("Decay",
+                   "Exponential decay coefficient, Hz; zero is a valid value for static scenarios",
+                   DoubleValue (1.0),
+                   MakeDoubleAccessor (&ThompsonSamplingWifiManager::m_decay),
+                   MakeDoubleChecker<double> (0.0))
+    .AddTraceSource ("Rate",
+                     "Traced value for rate changes (b/s)",
+                     MakeTraceSourceAccessor (&ThompsonSamplingWifiManager::m_currentRate),
+                     "ns3::TracedValueCallback::Uint64")
+  ;
+  return tid;
+}
+
+ThompsonSamplingWifiManager::ThompsonSamplingWifiManager ()
+  : m_currentRate{0}
+{
+  NS_LOG_FUNCTION (this);
+
+  m_gammaRandomVariable = CreateObject<GammaRandomVariable> ();
+}
+
+ThompsonSamplingWifiManager::~ThompsonSamplingWifiManager ()
+{
+  NS_LOG_FUNCTION (this);
+}
+
+void
+ThompsonSamplingWifiManager::SetupPhy (const Ptr<WifiPhy> phy)
+{
+  NS_LOG_FUNCTION (this << phy);
+  WifiRemoteStationManager::SetupPhy (phy);
+}
+
+WifiRemoteStation *
+ThompsonSamplingWifiManager::DoCreateStation () const
+{
+  NS_LOG_FUNCTION (this);
+  ThompsonSamplingWifiRemoteStation *station = new ThompsonSamplingWifiRemoteStation ();
+  station->m_nextMode = 0;
+  station->m_lastMode = 0;
+  return station;
+}
+
+void
+ThompsonSamplingWifiManager::InitializeStation (WifiRemoteStation *st) const
+{
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+  if (!station->m_mcsStats.empty ())
+    {
+      return;
+    }
+
+  // Add HT, VHT or HE MCSes
+  for (const auto &mode : GetPhy ()->GetMcsList ())
+    {
+      for (uint16_t j = 20; j <= GetPhy ()->GetChannelWidth (); j *= 2)
+        {
+          WifiModulationClass modulationClass = WIFI_MOD_CLASS_HT;
+          if (GetVhtSupported ())
+            {
+              modulationClass = WIFI_MOD_CLASS_VHT;
+            }
+          if (GetHeSupported ())
+            {
+              modulationClass = WIFI_MOD_CLASS_HE;
+            }
+          if (mode.GetModulationClass () == modulationClass)
+            {
+              for (uint8_t k = 1; k <= GetPhy ()->GetMaxSupportedTxSpatialStreams (); k++)
+                {
+                  if (mode.IsAllowed (j, k))
+                    {
+                      RateStats stats;
+                      stats.mode = mode;
+                      stats.channelWidth = j;
+                      stats.nss = k;
+
+                      station->m_mcsStats.push_back (stats);
+                    }
+                }
+            }
+        }
+    }
+
+  if (station->m_mcsStats.empty ())
+    {
+      // Add legacy non-HT modes.
+      for (uint8_t i = 0; i < GetNSupported (station); i++)
+        {
+          RateStats stats;
+          stats.mode = GetSupported (station, i);
+          if (stats.mode.GetModulationClass () == WIFI_MOD_CLASS_DSSS
+              || stats.mode.GetModulationClass () == WIFI_MOD_CLASS_HR_DSSS)
+            {
+              stats.channelWidth = 22;
+            }
+          else
+            {
+              stats.channelWidth = 20;
+            }
+          stats.nss = 1;
+          station->m_mcsStats.push_back (stats);
+        }
+    }
+
+  NS_ASSERT_MSG (!station->m_mcsStats.empty (), "No usable MCS found");
+
+  UpdateNextMode (st);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportRxOk (WifiRemoteStation *station, double rxSnr, WifiMode txMode)
+{
+  NS_LOG_FUNCTION (this << station << rxSnr << txMode);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportRtsFailed (WifiRemoteStation *station)
+{
+  NS_LOG_FUNCTION (this << station);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportDataFailed (WifiRemoteStation *st)
+{
+  NS_LOG_FUNCTION (this << st);
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+  Decay (st, station->m_lastMode);
+  station->m_mcsStats.at (station->m_lastMode).fails++;
+  UpdateNextMode (st);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportRtsOk (WifiRemoteStation *st, double ctsSnr, WifiMode ctsMode,
+                                     double rtsSnr)
+{
+  NS_LOG_FUNCTION (this << st << ctsSnr << ctsMode.GetUniqueName () << rtsSnr);
+}
+
+void
+ThompsonSamplingWifiManager::UpdateNextMode (WifiRemoteStation *st) const
+{
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+
+  double maxThroughput = 0.0;
+  double frameSuccessRate = 1.0;
+
+  NS_ASSERT (!station->m_mcsStats.empty ());
+
+  // Use the most robust MCS if frameSuccessRate is 0 for all MCS.
+  station->m_nextMode = 0;
+
+  for (uint32_t i = 0; i < station->m_mcsStats.size (); i++)
+    {
+      Decay (st, i);
+      const WifiMode mode{station->m_mcsStats.at (i).mode};
+
+      uint16_t guardInterval = GetModeGuardInterval (st, mode);
+      double rate = mode.GetDataRate (station->m_mcsStats.at (i).channelWidth,
+                                      guardInterval,
+                                      station->m_mcsStats.at (i).nss);
+
+      // Thompson sampling
+      frameSuccessRate = SampleBetaVariable (1.0 + station->m_mcsStats.at (i).success,
+                                             1.0 + station->m_mcsStats.at (i).fails);
+      NS_LOG_DEBUG ("Draw"
+                    << " success=" << station->m_mcsStats.at (i).success
+                    << " fails=" << station->m_mcsStats.at (i).fails
+                    << " frameSuccessRate=" << frameSuccessRate
+                    << " mode=" << mode);
+      if (frameSuccessRate * rate > maxThroughput)
+        {
+          maxThroughput = frameSuccessRate * rate;
+          station->m_nextMode = i;
+        }
+    }
+}
+
+void
+ThompsonSamplingWifiManager::DoReportDataOk (WifiRemoteStation *st, double ackSnr, WifiMode ackMode,
+                                      double dataSnr, uint16_t dataChannelWidth, uint8_t dataNss)
+{
+  NS_LOG_FUNCTION (this << st << ackSnr << ackMode.GetUniqueName () << dataSnr);
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+  Decay (st, station->m_lastMode);
+  station->m_mcsStats.at (station->m_lastMode).success++;
+  UpdateNextMode (st);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportAmpduTxStatus (WifiRemoteStation *st, uint16_t nSuccessfulMpdus,
+                                             uint16_t nFailedMpdus, double rxSnr, double dataSnr,
+                                             uint16_t dataChannelWidth, uint8_t dataNss)
+{
+  NS_LOG_FUNCTION (this << st << nSuccessfulMpdus << nFailedMpdus << rxSnr << dataSnr);
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+
+  Decay (st, station->m_lastMode);
+  station->m_mcsStats.at (station->m_lastMode).success += nSuccessfulMpdus;
+  station->m_mcsStats.at (station->m_lastMode).fails += nFailedMpdus;
+
+  UpdateNextMode (st);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportFinalRtsFailed (WifiRemoteStation *station)
+{
+  NS_LOG_FUNCTION (this << station);
+}
+
+void
+ThompsonSamplingWifiManager::DoReportFinalDataFailed (WifiRemoteStation *station)
+{
+  NS_LOG_FUNCTION (this << station);
+}
+
+uint16_t
+ThompsonSamplingWifiManager::GetModeGuardInterval (WifiRemoteStation *st, WifiMode mode) const
+{
+  if (mode.GetModulationClass () == WIFI_MOD_CLASS_HE)
+    {
+      return std::max (GetGuardInterval (st), GetGuardInterval ());
+    }
+  else if ((mode.GetModulationClass () == WIFI_MOD_CLASS_HT) ||
+           (mode.GetModulationClass () == WIFI_MOD_CLASS_VHT))
+    {
+      return std::max<uint16_t> (GetShortGuardIntervalSupported (st) ? 400 : 800,
+                                 GetShortGuardIntervalSupported () ? 400 : 800);
+    }
+  else
+    {
+      return 800;
+    }
+}
+
+WifiTxVector
+ThompsonSamplingWifiManager::DoGetDataTxVector (WifiRemoteStation *st)
+{
+  NS_LOG_FUNCTION (this << st);
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+
+  auto &stats = station->m_mcsStats.at (station->m_nextMode);
+  WifiMode mode = stats.mode;
+  uint16_t channelWidth = std::min (stats.channelWidth, GetPhy ()->GetChannelWidth ());
+  uint8_t nss = stats.nss;
+  uint16_t guardInterval = GetModeGuardInterval (st, mode);
+
+  station->m_lastMode = station->m_nextMode;
+
+  NS_LOG_DEBUG ("Using"
+                << " mode=" << mode
+                << " channelWidth=" << channelWidth
+                << " nss=" << +nss
+                << " guardInterval=" << guardInterval);
+
+  uint64_t rate = mode.GetDataRate (channelWidth, guardInterval, nss);
+  if (m_currentRate != rate)
+    {
+      NS_LOG_DEBUG ("New datarate: " << rate);
+      m_currentRate = rate;
+    }
+
+  return WifiTxVector (
+      mode,
+      GetDefaultTxPowerLevel (),
+      GetPreambleForTransmission (mode.GetModulationClass (),
+                                  GetShortPreambleEnabled ()),
+      GetModeGuardInterval (st, mode),
+      GetNumberOfAntennas (),
+      nss,
+      0, // NESS
+      GetChannelWidthForTransmission (mode, channelWidth),
+      GetAggregation (station),
+      false);
+}
+
+WifiTxVector
+ThompsonSamplingWifiManager::DoGetRtsTxVector (WifiRemoteStation *st)
+{
+  NS_LOG_FUNCTION (this << st);
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+
+  // Use the most robust MCS for the control channel.
+  auto &stats = station->m_mcsStats.at (0);
+  WifiMode mode = stats.mode;
+  uint16_t channelWidth = std::min (stats.channelWidth, GetPhy ()->GetChannelWidth ());
+  uint8_t nss = stats.nss;
+
+  // Make sure control frames are sent using 1 spatial stream.
+  NS_ASSERT (nss == 1);
+
+  return WifiTxVector (
+      mode, GetDefaultTxPowerLevel (),
+      GetPreambleForTransmission (mode.GetModulationClass (), GetShortPreambleEnabled ()),
+      GetModeGuardInterval (st, mode),
+      GetNumberOfAntennas (),
+      nss,
+      0, // NESS
+      GetChannelWidthForTransmission (mode, channelWidth),
+      GetAggregation (station),
+      false);
+}
+
+double
+ThompsonSamplingWifiManager::SampleBetaVariable (uint64_t alpha, uint64_t beta) const
+{
+  double X = m_gammaRandomVariable->GetValue (alpha, 1.0);
+  double Y = m_gammaRandomVariable->GetValue (beta, 1.0);
+  return X / (X + Y);
+}
+
+void
+ThompsonSamplingWifiManager::Decay (WifiRemoteStation *st, size_t i) const
+{
+  NS_LOG_FUNCTION (this << st << i);
+  InitializeStation (st);
+  auto station = static_cast<ThompsonSamplingWifiRemoteStation *> (st);
+
+  Time now = Simulator::Now ();
+  auto &stats = station->m_mcsStats.at (i);
+  if (now > stats.lastDecay)
+    {
+      const double coefficient =
+          std::exp (m_decay * (stats.lastDecay - now).GetSeconds ());
+
+      stats.success *= coefficient;
+      stats.fails *= coefficient;
+      stats.lastDecay = now;
+    }
+}
+
+int64_t
+ThompsonSamplingWifiManager::AssignStreams (int64_t stream)
+{
+  NS_LOG_FUNCTION (this << stream);
+  m_gammaRandomVariable->SetStream (stream);
+  return 1;
+}
+
+} //namespace ns3
diff --git a/src/wifi/model/rate-control/thompson-sampling-wifi-manager.h b/src/wifi/model/rate-control/thompson-sampling-wifi-manager.h
new file mode 100644
index 000000000..88541ac62
--- /dev/null
+++ b/src/wifi/model/rate-control/thompson-sampling-wifi-manager.h
@@ -0,0 +1,135 @@
+/* -*- Mode:C++; c-file-style:"gnu"; indent-tabs-mode:nil; -*- */
+/*
+ * Copyright (c) 2021 IITP RAS
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License version 2 as
+ * published by the Free Software Foundation;
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+ *
+ * Author: Alexander Krotov <krotov@iitp.ru>
+ */
+
+#ifndef THOMPSON_SAMPLING_WIFI_MANAGER_H
+#define THOMPSON_SAMPLING_WIFI_MANAGER_H
+
+#include "ns3/random-variable-stream.h"
+
+#include "ns3/wifi-remote-station-manager.h"
+
+namespace ns3 {
+
+/**
+ * \brief Thompson Sampling rate control algorithm
+ * \ingroup wifi
+ *
+ * This class implements Thompson Sampling rate control algorithm.
+ *
+ * It was implemented for use as a baseline in
+ * https://doi.org/10.1109/ACCESS.2020.3023552
+ */
+class ThompsonSamplingWifiManager : public WifiRemoteStationManager
+{
+public:
+  /**
+   * \brief Get the type ID.
+   * \return the object TypeId
+   */
+  static TypeId GetTypeId (void);
+  ThompsonSamplingWifiManager ();
+  virtual ~ThompsonSamplingWifiManager ();
+
+  void SetupPhy (const Ptr<WifiPhy> phy) override;
+
+  /**
+   * Assign a fixed random variable stream number to the random variables
+   * used by this model.  Return the number of streams (possibly zero) that
+   * have been assigned.
+   *
+   * \param stream first stream index to use
+   * \return the number of stream indices assigned by this model
+   */
+  int64_t AssignStreams (int64_t stream);
+
+private:
+  WifiRemoteStation *DoCreateStation () const override;
+  void DoReportRxOk (WifiRemoteStation *station,
+                     double rxSnr, WifiMode txMode) override;
+  void DoReportRtsFailed (WifiRemoteStation *station) override;
+  void DoReportDataFailed (WifiRemoteStation *station) override;
+  void DoReportRtsOk (WifiRemoteStation *station,
+                      double ctsSnr, WifiMode ctsMode, double rtsSnr) override;
+  void DoReportDataOk (WifiRemoteStation *station,
+                       double ackSnr, WifiMode ackMode, double dataSnr,
+                       uint16_t dataChannelWidth, uint8_t dataNss) override;
+  void DoReportAmpduTxStatus (WifiRemoteStation *station,
+                              uint16_t nSuccessfulMpdus, uint16_t nFailedMpdus,
+                              double rxSnr, double dataSnr, uint16_t dataChannelWidth, uint8_t dataNss) override;
+  void DoReportFinalRtsFailed (WifiRemoteStation *station) override;
+  void DoReportFinalDataFailed (WifiRemoteStation *station) override;
+  WifiTxVector DoGetDataTxVector (WifiRemoteStation *station) override;
+  WifiTxVector DoGetRtsTxVector (WifiRemoteStation *station) override;
+
+  /**
+   * Initializes station rate tables. If station is already initialized,
+   * nothing is done.
+   *
+   * \param station Station which should be initialized.
+   */
+  void InitializeStation (WifiRemoteStation *station) const;
+
+  /**
+   * Draws a new MCS and related parameters to try next time for this
+   * station.
+   *
+   * This method should only be called between TXOPs to avoid sending
+   * multiple frames using different modes. Otherwise it is impossible
+   * to tell which mode was used for succeeded/failed frame when
+   * feedback is received.
+   *
+   * \param station Station for which a new mode should be drawn.
+   */
+  void UpdateNextMode (WifiRemoteStation *station) const;
+
+  /**
+   * Applies exponential decay to MCS statistics.
+   *
+   * \param st Remote STA for which MCS statistics is to be updated.
+   * \param i MCS index.
+   */
+  void Decay (WifiRemoteStation *st, size_t i) const;
+
+  /**
+   * Returns guard interval in nanoseconds for the given mode.
+   *
+   * \param st Remote STA.
+   * \param mode The WifiMode.
+   */
+  uint16_t GetModeGuardInterval (WifiRemoteStation *st, WifiMode mode) const;
+
+  /**
+   * Sample beta random variable with given parameters
+   * \param alpha first parameter of beta distribution
+   * \param beta second parameter of beta distribution
+   * \return beta random variable sample
+   */
+  double SampleBetaVariable (uint64_t alpha, uint64_t beta) const;
+
+  Ptr<GammaRandomVariable> m_gammaRandomVariable; //< Variable used to sample beta-distributed random variables
+
+  double m_decay; //< Exponential decay coefficient, Hz
+
+  TracedValue<uint64_t> m_currentRate; //!< Trace rate changes
+};
+
+} //namespace ns3
+
+#endif /* THOMPSON_SAMPLING_WIFI_MANAGER_H */
diff --git a/src/wifi/wscript b/src/wifi/wscript
index cb87de87a..19a633baa 100644
--- a/src/wifi/wscript
+++ b/src/wifi/wscript
@@ -83,6 +83,7 @@ def build(bld):
         'model/rate-control/parf-wifi-manager.cc',
         'model/rate-control/aparf-wifi-manager.cc',
         'model/rate-control/rrpaa-wifi-manager.cc',
+        'model/rate-control/thompson-sampling-wifi-manager.cc',
         'model/ampdu-subframe-header.cc',
         'model/mpdu-aggregator.cc',
         'model/ampdu-tag.cc',
@@ -193,6 +194,7 @@ def build(bld):
         'model/rate-control/cara-wifi-manager.h',
         'model/rate-control/minstrel-wifi-manager.h',
         'model/rate-control/minstrel-ht-wifi-manager.h',
+        'model/rate-control/thompson-sampling-wifi-manager.h',
         'model/wifi-mac.h',
         'model/regular-wifi-mac.h',
         'model/supported-rates.h',