PyMAB

PyMAB is a small Python library for reproducible multi-armed bandit experiments. It separates environments, policies, simulations, metrics, and plotting so experiments are easier to test, compare, and extend.

What changed in v1

• Python 3.11+ and a clean typed API.
• Explicit random seeds through numpy.random.Generator.
• Environments own true arm values and reward sampling.
• Policies only select actions and update from observed rewards.
• Regret is computed from expected rewards, with realized rewards kept separate.
• Plotting dependencies are optional via pymab[plot].
• Pandas analysis helpers are optional via pymab[analysis].
• Benchmarks can compare policies across repeated seeds with confidence intervals.

Install

pip install pymab

For local development:

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev,plot]"
python -m unittest discover -v

With uv, use the Makefile targets used by CI:

make sync
make format
make lint
make security
make test

make format checks formatting. Use make format-fix to rewrite formatting locally.

The Makefile tries uv run ... first. If your local uv binary fails before it can run a tool, the targets fall back to executables installed in .venv; run make sync once to create/populate that environment.

Basic Example

import numpy as np

from pymab.environments import BanditEnvironment
from pymab.policies import EpsilonGreedyPolicy, UCBPolicy
from pymab.simulation import Experiment, ExperimentConfig

environment = BanditEnvironment(q_values=np.array([0.1, 0.4, 0.8]))
policies = [
    EpsilonGreedyPolicy(n_arms=3, epsilon=0.1),
    UCBPolicy(n_arms=3, c=2.0),
]

result = Experiment(
    environment=environment,
    policies=policies,
    config=ExperimentConfig(n_episodes=200, n_steps=500, seed=42),
).run()

print(result.average_reward_by_step[-1])
print(result.cumulative_regret[-1])

Compare Policies

import numpy as np

from pymab import compare
from pymab.environments import BanditEnvironment
from pymab.policies import RandomPolicy, ThompsonSamplingPolicy, UCBPolicy

benchmark = compare(
    [
        RandomPolicy(n_arms=3),
        UCBPolicy(n_arms=3),
        ThompsonSamplingPolicy(n_arms=3),
    ],
    environment=BanditEnvironment(q_values=np.array([0.1, 0.3, 0.8])),
    n_episodes=100,
    n_steps=500,
    seeds=(1, 2, 3, 4, 5),
)

print(benchmark.best_policy)
print(benchmark.summary())

Install pymab[analysis] to convert results into pandas DataFrames:

result_frame = benchmark.combined.to_pandas()
summary_frame = benchmark.to_pandas()

Persist reproducible result arrays with:

benchmark.combined.save_npz("results/benchmark.npz")

Non-Stationary Environments

import numpy as np

from pymab.environments import BanditEnvironment, GradualDrift

environment = BanditEnvironment(
    q_values=np.array([0.2, 0.5, 0.7]),
    dynamics=GradualDrift(change_rate=0.01),
)

Built-in dynamics:

• StationaryDynamics
• GradualDrift
• AbruptShift
• RandomArmSwap

Policies

Classic bandits:

• GreedyPolicy
• RandomPolicy
• EpsilonGreedyPolicy
• DecayingEpsilonGreedyPolicy
• SoftmaxPolicy
• GradientBanditPolicy
• UCBPolicy
• KLUCBPolicy
• MOSSPolicy
• SlidingWindowUCBPolicy
• DiscountedUCBPolicy
• CUSUMUCBPolicy
• PageHinkleyUCBPolicy
• BernoulliThompsonSamplingPolicy
• GaussianThompsonSamplingPolicy
• SlidingWindowBernoulliThompsonSamplingPolicy
• DiscountedBernoulliThompsonSamplingPolicy
• BernoulliBayesianUCBPolicy
• GaussianBayesianUCBPolicy
• EXP3Policy
• SuccessiveEliminationPolicy
• MedianEliminationPolicy

Contextual bandits:

• LinearEpsilonGreedyPolicy
• LinUCBPolicy
• LinearThompsonSamplingPolicy
• LogisticContextualBanditPolicy

Contextual Example

import numpy as np

from pymab.environments import LinearContextualEnvironment
from pymab.policies import LinUCBPolicy
from pymab.simulation import Experiment, ExperimentConfig


def context_provider(rng: np.random.Generator) -> np.ndarray:
    return np.array([[1.0, 0.0], [0.0, 1.0]])


environment = LinearContextualEnvironment(
    theta=np.array([[1.0, 0.0], [0.0, 1.0]]),
    context_provider=context_provider,
)

result = Experiment(
    environment=environment,
    policies=[LinUCBPolicy(n_arms=2, n_features=2)],
    config=ExperimentConfig(n_episodes=100, n_steps=200, seed=7),
).run()

Plotting

from pathlib import Path

from pymab.plotting import plot_average_reward, plot_cumulative_regret

plot_average_reward(result, output_path=Path("results/average_reward.html"))
plot_cumulative_regret(result, output_path=Path("results/regret.html"))

Install plotting extras first:

pip install "pymab[plot]"

Migration Notes

The old Game API remains as a deprecated compatibility wrapper. New code should use BanditEnvironment, ExperimentConfig, and Experiment directly. The old pymab.reward_distribution import path also remains available, but the new module is pymab.distributions.