Closing the loop after PolicyLearner: batched Thompson sampling on the treatment-effect posterior (scoping, following #833)

Open
  • Is something wrong?
#1,058 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
20/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
numpy, python
Domain
data, machine-learning

Research direction

This is a design proposal that asks maintainers to decide whether a batched Thompson sampler belongs in causalml/optimize/bandit.py before any code is written. Read the PolicyLearner code in causalml/optimize and the "Choosing an Estimator" docs page to see how randomized, known-propensity data is handled. Done means a maintainer agrees on placement and scope, then a sampler and tests that use only numpy and scipy.

Written by the indexing model from the issue text.

Description

When #833 was closed, you suggested agreeing first on where multi-armed bandits fit relative to causalml.optimize. Here is a narrow proposal, and the question of whether it fits.

Where it would sit. CausalML today takes one dataset, estimates the effect, validates it, and decides who to treat (PolicyLearner, value optimization). The proposal adds the step after that decision: run it as the next randomized batch, update the effect estimate, and decide again. It is one method, not a general bandit collection, and the loop is built around a treatment-effect estimate.

The core: an effect posterior carried from batch to batch.

  • Each batch is a randomized experiment with known, fixed assignment shares.
  • Within a batch, the method fits a logistic regression with the batch's own intercept and each arm's log odds ratio against a control arm. Only the effects are carried forward: their posterior becomes the prior for the next batch, and the intercept is refit every batch. Anything that moves all arms together within a batch, such as a promotion or a weekday, goes into the intercept instead of the effects.
  • The next batch's shares come from Thompson sampling on the effect posterior.
  • Because the shares depend only on the batch, the batch intercept stratifies on exactly what the assignment depended on. It is the comparison a propensity-stratified estimate makes, done through the design.

The logs are ordinary causalml input. Every unit was randomized with a known probability, its batch's share for the arm it received (p = shares[treatment]). The accumulated log (X, treatment, y, p, batch) can therefore be analysed afterwards by the existing learners with p passed in, as the "Choosing an Estimator" page describes for randomized data with a known assignment probability. The bandit handles allocation online; CATE estimation, validation and policy_risk stay with the existing tools.

Scale. Uplift in causalml is mostly a risk difference, while this method works on log odds. Within a segment both rank the arms the same way, so allocation is unaffected; reporting a risk difference needs the segment's base rate.

Suggested first PR, following your note on #833: causalml/optimize/bandit.py with

  • the batched effect-posterior sampler for a binary outcome and discrete arms, without covariates;
  • a Beta-Bernoulli Thompson sampler as the baseline;
  • tests, formatted with black, using only numpy and scipy.

I'd suggest marking it experimental so that it stays outside the v1.0 API freeze. Segments (each segment's effect shrunk toward the common effect by an amount estimated from the data), simulation utilities and a notebook would follow as separate PRs, and only if the first one fits.

A rough shape, with names open to change:

sampler = BatchedEffectThompsonSampler(n_arms=3, control_name=0, random_state=0)
shares = sampler.allocate()        # assignment shares for the next batch
sampler.update(treatment=t, y=y)   # outcomes observed in that batch

Background: the method is described in arXiv:2609.19709, and I maintain an MIT implementation (orts on PyPI). I would write the causalml version to the library's conventions rather than add a dependency. If this doesn't fit, or belongs somewhere other than optimize, I'd rather know before writing code.

Dominant language
Python
Stars
6k
Forks
889
Avg merge
11h 53m
Merged PRs (30d)
1

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from uber/causalml

All issues in uber/causalml

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.