Thompson Sampling with Stochastic Approximation (TS-SA)

This repository contains implementation of the Thompson Sampling with Stochastic Approximation (TS-SA) algorithm. The codebase has been simplified to include the core TS-SA implementation to run experiments.

Repository Structure

TS-SA/
├── config/                    # Active JSON configurations (see `config/ts_sa.json`)
├── examples/                  # Minimal runnable examples
├── main_ts_sa.sh              # Launcher script for TS-SA experiments
├── run_ts.py                  # Core script to execute TS-SA with specified parameters
├── model/                     # Core implementation of TS-SA and utilities
│   ├── __init__.py            # Module initializer (exports TS, TS_SA)
│   ├── algorithms.py          # Minimal TS-SA implementation and helpers (active)
│   ├── flags.py               # Command-line argument definitions and flag parsing
│   ├── reward.py              # Reward generation and synthetic bandit environments
│   └── utils.py               # Helper functions for printing
├── logs/                      # Output directory for experiment results
├── tests/                     # Minimal tests to verify correctness (pytest)
├── environment.yml            # Conda environment specification with required dependencies
└── README.md                  # Project documentation and usage instructions

Prerequisites

• We use Python 3.10.0.
• To set up a Python environment using conda:

conda env create -f environment.yml
conda activate thompson

Some useful packages included in environment.yml:

• numpy==1.26.0
• scipy==1.15.2
• scikit-learn==1.6.1
• tqdm==4.66.1
• matplotlib==3.5.0
• optuna==4.2.1
• pyarrow==19.0.1
• pillow==10.0.1

Check environment.yml for the full list and exact pinned builds.

Usage

The trimmed repo supports running TS-SA via the run_ts.py script. A minimal example is provided under examples/run_ts_sa.py.

Run the example:

python run_ts_sa.py

Or run a short experiment using the launcher script (it reads config/ts_sa.json by default):

bash main_ts_sa.sh

The launcher currently executes the sa algorithm (TS-SA).

Output

All experiment outputs and saved traces are written to the logs/ directory.

References

If you find this repository useful, please cite the following work:

@article{wang2025rethinking,
  title={Rethinking Langevin Thompson Sampling from A Stochastic Approximation Perspective},
  author={Wang, Weixin and Zheng, Haoyang and Lin, Guang and Deng, Wei and Xu, Pan},
  journal={arXiv preprint arXiv:2510.05023},
  year={2025}
}