contextual-bandit-dispatch

Offline contextual bandit for dispatch optimization on Graph Data for supply chain of pallets.

Overview

This project learns a dispatching policy from logged decisions and applies it to recommend a ServiceCenter (SC) for each open WorkOrder (WO). It runs in two layers:

• A Python pipeline that pulls historical (context, action, reward) data from TigerGraph, trains five models, and writes per-candidate predictions back to TigerGraph as edge attributes.
• A set of GSQL queries that consume those predictions to produce business-readable recommendations with tunable cost weights.

The clean separation means data scientists can retrain models without touching business logic, and operators can re-weight cost penalties without retraining.

What the system does

Given:

• A history of WorkOrder → ServiceCenter assignments with realized outcomes (was_late, had_damage, had_overflow, costs)
• A set of currently-open WorkOrders and their candidate ServiceCenters

The system produces, for every (open WO, candidate SC) pair:

• A predicted probability of late delivery, damage, and overflow
• A learned policy score
• A behavior propensity (how likely the historical policy was to make this choice given the WO context)
• A SHAP-based reasoning string explaining the prediction
• An expected total cost combining direct cost with risk-adjusted penalties

A GSQL query then picks the cost-minimizing SC per open WO using business-tunable penalty weights, and exposes a recommendation table any downstream consumer can query.

Architecture

┌─────────────────┐         ┌──────────────────────────────────┐
│   TigerGraph    │         │       Python pipeline            │
│                 │  pull   │                                  │
│  Logged WOs     │────────▶│  Encoder  (one-hot + scaling)    │
│  Open WOs       │         │  Propensity  μ̂(a|x)              │
│  Outcomes       │         │  Reward      r̂(x, a)             │
│                 │         │  Learner     π_θ  (PyTorch MLP)  │
│                 │         │  Failure×3   p̂_late, p̂_dmg, p̂_ov │
│                 │         │  SHAP attribution                │
│                 │  write  │                                  │
│  Edge predicts  │◀────────│  Writeback                       │
└────────┬────────┘         └──────────────────────────────────┘
         │
         │ run query
         ▼
┌─────────────────────────────────────────────────────────────┐
│            GSQL business layer                              │
│                                                             │
│  recommend_sc_for_workorders(δ, λ, ω, n)                    │
│    expected_total_cost = direct_cost                        │
│                        + δ · p_late     · qty               │
│                        + λ · p_damage   · qty               │
│                        + ω · p_overflow · qty               │
│    argmin per WO → write to vertex; emit sorted table       │
└─────────────────────────────────────────────────────────────┘

The math

The system implements an offline contextual bandit. Three core estimators feed a Doubly-Robust (DR) policy value estimator (Dudík, Langford & Li, 2011):

Propensity model μ̂(a | x). A multinomial logistic regression on the WO-context features only — excluding action-side features to avoid feature/target leakage. Outputs the probability that the historical (logging) policy would have chosen action a given context x.

Reward model r̂(x, a). An XGBoost regressor estimating expected reward given context and action. Used as both the "direct method" predictor and as a baseline for the IPS correction.

Learner π_θ. An MLP scoring (context, action) feature vectors. Trained to maximize the soft (probabilistic) version of V̂_DR:

V̂_soft(π_θ) = (1/N) Σ_i [  Σ_j π_θ(c_j|x_i) · r̂(x_i, c_j)
                          + π_θ(a_i|x_i) · (r_i − r̂(x_i, a_i))
                                          / μ̂(a_i|x_i) ]

The first term is the direct method (use r̂); the second is the IPS correction (reweight observed rewards by 1/μ̂). The estimator is doubly robust: it's unbiased if either r̂ or μ̂ is correctly specified.

Failure models p̂_k(x, a). Three binary XGBoost classifiers, one per failure mode (late, damage, overflow). These are what surface as risk probabilities in the recommendation table.

Reward construction. Realized rewards are negative cost:

r_i = −( delay_penalty    · was_late     · qty
      + damage_penalty   · had_damage   · qty
      + overflow_penalty · had_overflow · qty
      + realized_direct_cost )

Scaled by 1/1000 before training so values land in a numerically friendly range and downstream models don't overfit to noise on the high-magnitude tail.

Installation

Requires Python 3.14+ and a TigerGraph instance with the required queries installed (see "TigerGraph queries" below).

git clone <repo>
cd contextual-bandit-dispatch

python3.14 -m venv .venv
source .venv/bin/activate

pip install --upgrade pip
pip install -e ".[dev]"

Configuration

Copy .env.example to .env and fill in:

HOST=https://your-instance.i.tgcloud.io
GRAPHNAME=YOUR_GRAPH_NAME
SECRET=<your secret>

These are read via pydantic-settings at startup. Missing values raise ValidationError with a clear message naming the missing field.

Usage

Run the Python pipeline

python -m contextual_bandit_dispatch.cli

This:

1. Pulls logged candidate decisions from TigerGraph.
2. Trains every model in the bandit (Propensity, Reward, Learner, three Failure classifiers).
3. Compares the learner against three heuristic baselines via V̂_DR.
4. Scores every open candidate edge and computes failure probabilities plus SHAP reasoning for each.
5. Writes all predictions back to TigerGraph edges.

Expected output: stage-level progress logs, a policy comparison table, and a confirmation count of edges pushed.

Produce recommendations from TigerGraph

After the pipeline finishes, per-candidate predictions live on WORK_ORDER_FULFILLMENT_CANDIDATE edges. To turn those into business-readable recommendations:

RUN QUERY recommend_sc_for_workorders(50.0, 200.0, 100.0, 1000)

The first three parameters are penalty weights for delay, damage, and overflow respectively. The fourth caps the result table size. The query:

• Computes expected_total_cost per candidate edge.
• Picks the argmin per WO.
• Writes recommended_sc_id, expected_total_cost, and a composed recommendation_explanation to each open WO vertex.
• Flags the winning candidate edge with is_recommended = 1.
• Returns a sorted table (one row per open WO) plus diagnostic counts.

Re-running with different penalty weights produces different recommendations without retraining anything in Python. This is the operator-facing knob.

TigerGraph queries

The project assumes four installed queries on the graph:

Query	Purpose
extract_logged_decisions_with_realized_rewards(partition_id, total_partitions)	Returns historical (context, action, reward) rows as typed tuples. Partition-aware for parallel pulls.
extract_open_wo_candidate_contexts(partition_id, total_partitions)	Returns (open WO, candidate SC) rows for scoring. Same partition contract.
writeback_candidate_predictions(...)	Receives per-candidate predictions from Python and writes them to edge attributes.
recommend_sc_for_workorders(δ, λ, ω, limit_n)	Cost-min recommendation engine. Reads edge predictions, applies penalty weights, writes recommendations to WO vertices, emits a sorted table.

GSQL source files live alongside the repo and are installed via the TG admin shell (gsql -g <graph> < query.gsql).

Development

# Type checking (strict)
basedpyright src/contextual_bandit_dispatch

# Lint + import sorting
ruff check src/

# Tests
pytest                           # unit tests
pytest -m integration            # tests requiring a live TG connection

The project enforces strict basedpyright; uses cast() at library boundaries rather than # pyright: ignore suppressions.

Algorithmic notes

A few design decisions worth knowing if you're extending the code:

Propensity uses context-only features. The Propensity model deliberately excludes sc_id and every action-side feature (sc_capacity, distance_km, sc_lifetime_*, etc.) — otherwise the model trivially identifies which action was chosen, producing degenerate near-{0, 1} propensities and unstable IPS corrections. The split is enforced by separate CONTEXT_CATEGORICAL_FEATURES and CONTEXT_NUMERIC_FEATURES tuples in tigergraph/headers.py.

Reward scaling. Raw rewards are in the thousands (penalties × pallet quantities). Dividing by 1000 before any model sees them keeps XGBoost from overfitting wildly and the PyTorch DR loss from producing absurd magnitudes. Scaling is uniform across all downstream models so policy ordering is preserved.

Propensity floor in evaluation. The DR estimator floors propensities at 0.01 to bound IPS variance. Without it, a single row with μ̂ ≈ 1e-4 could swing V̂_DR by orders of magnitude. The floor is in evaluation/value_estimator.py; lower it (or remove it) on production-grade data where extreme propensities carry real signal.

Single-threaded torch. The Learner forces torch.set_num_threads(1) because PyTorch's default OpenMP thread spawn is pathologically slow on macOS, and there's no parallelism benefit at this data size anyway.

Predict for all candidates, not just picks. workflow.recommend() scores every open candidate, not just the argmax per WO. This is what makes the cost-min query in GSQL meaningful — it has real predictions on every option, so penalty-weight changes produce real argmin changes across the candidate set.

Argmin discovery is two passes. The TG recommend_sc_for_workorders query splits argmin discovery (Step 1: MinAccum tracks running minimum) from detail capture (Step 2: filter by expected_total_cost == best_cost). This avoids a known GSQL race condition where conditional captures inside an ACCUM block can produce nondeterministic results.

License

MIT.