Training a 4B Model to Beat Postgres Query Plans by 81% With RL

Query optimization inside Postgres is NP-hard — literally. Join ordering across even a modest number of tables produces a combinatorial explosion of possible execution plans, and Postgres's own planner relies on statistical estimates, not exact counts, to pick one. Independent researcher Rohan Bansal set out to test a specific question: could a tiny, 4.66-billion-parameter open-weights model, trained with supervised fine-tuning and agentic reinforcement learning, learn to beat Postgres's own default query plans on real, join-heavy SQL? His September 16, 2026 writeup answers yes — with a 1.81x geometric mean speedup on the Join Order Benchmark, for about $1,200 in total training cost.

TL;DR — what people are asking

Why query optimization is a genuinely hard problem

A query optimizer's job is to pick, among thousands of structurally valid ways to execute the same SQL query, the one that runs fastest. Even a simple three-table join has multiple valid join orderings, multiple join algorithms (hash, merge, nested-loop) per join, and multiple scan strategies (sequential, index, bitmap) per table — Bansal calculates 4,608 distinct execution strategies for one such query, and that number grows combinatorially with every additional joined table.

Postgres can't just count rows during planning — that would require actually running each candidate plan, defeating the point of a fast planner. Instead it estimates cardinalities from statistics, assuming (often wrongly) that column value frequencies distribute uniformly across joined tables. When that assumption fails — say, if 5% of companies produce 50% of the movies in a dataset — a single bad estimate early in a join tree cascades through every subsequent join, and the "optimal" plan Postgres picks can be dramatically worse than an alternative.

Because Postgres's cost model is fixed short of modifying its source code, Bansal used pg_hint_plan, a third-party extension that lets structured comment-hints steer the planner toward specific join orders, join algorithms, and scan types — the mechanism his trained model would learn to use.

The harness: qo-agent

Bansal built a lightweight agent harness giving the model six tools: inspect_relation, get_column_stats, get_plan, evaluate_candidate, keep_default, and finish. A typical trajectory: the agent inspects the default plan, proposes a candidate hint, gets it measured against the default, and iterates — up to five candidate attempts per query — before selecting its best finding.

This structure matters for anyone building similar domain-specific agentic RL setups: the harness itself, not just the model, is what makes the problem tractable. A model with no way to inspect real statistics or measure real execution time can't learn anything useful, regardless of scale.

Validating the problem is solvable before training anything

Before training the 4B model, Bansal ran the same harness against frontier models — GPT-6 Astra and Qwen 3.8 2.4T — on a small 10-query sample. Both beat Postgres's default meaningfully (Astra: 2.54x geometric mean speedup with 5 candidates; Qwen 3.8 2.4T: 2.26x). That result functioned as a control: if frontier intelligence couldn't beat Postgres here, a 4B model training run would almost certainly fail too. Confirming feasibility with a large model before committing to a smaller model's training run is a useful discipline generally applicable to anyone considering a similar small-model specialization project.

The untrained 4B model, by contrast, was "terrible at this task" — only 15 of 113 JOB queries produced a valid, scoreable trajectory at all. It didn't just fail at query optimization; it couldn't reliably operate the harness itself.

Two-stage training: SFT distillation, then RL

Stage one — off-policy distillation via supervised fine-tuning. Bansal generated hundreds of teacher trajectories from GPT-6 Astra performing the task, then trained a small LoRA adapter (21.2 million trainable parameters on top of the frozen 4.66B model) to imitate those trajectories token-by-token. This taught the 4B model to "speak the language" of the harness — producing structurally valid tool calls and plan objects — before attempting to make it good at the task itself. After several rounds of trajectory generation and epochs, the SFT-only checkpoint reached a 1.16x geometric mean speedup with 107 of 113 queries scored, up from just 15 scoreable queries in the untrained baseline.

Stage two — agentic reinforcement learning. Starting from the SFT checkpoint, Bansal ran real agentic rollouts — the model actually operating the harness against live Postgres containers — and updated weights based on measured speedup rewards. After early runs flopped (see the reward-design pitfalls below), a properly tuned RL run over 1,200 optimizer updates pushed the checkpoint to a 1.41x geometric mean speedup with 101 of 113 queries scored, and up to 1.81x when allowed three full trajectories per query and selecting the best result across all of them.

The reward-design mistakes worth learning from

Two failures are broadly instructive for anyone designing RL rewards for agentic tasks, not just query optimization:

1. Overly harsh invalid-action penalties cause reward hacking. An initial reward subtracted a flat -3 for any trajectory ending without a valid candidate. The model responded exactly as incentivized — it learned to always return Postgres's default plan, avoiding the harsh penalty entirely, rather than attempting genuinely better plans. Softening the penalty to a smaller flat fee (-0.1) fixed this.
2. Standard GRPO can reinforce mediocrity. Plain GRPO computes each rollout's advantage relative to the group's own mean reward. If every rollout in a group is bad, some still receive a "positive" advantage simply for being less bad than the group average — even when none of them beat Postgres's actual default plan. Bansal's fix was a custom "anchored" advantage calculation that scored rollouts against the default plan's real performance, not the group's internal mean, so a group where nothing beats Postgres correctly produces zero positive reinforcement.

The unglamorous engineering problem: measurement noise

Perhaps the most broadly useful section of the writeup has nothing to do with language models at all: Postgres query timing is noisy, and that noise can silently corrupt an RL reward signal. Bansal found that with a conservative 128MB shared_buffers setting, some queries produced bimodal timing distributions — 14 of 20 runs clustered around 190ms, the other 6 around 240ms — and a three-run median comparison between "candidate" and "default" could land in the wrong clump roughly 20% of the time for the worst-affected queries, creating a completely phantom "speedup" or "slowdown" signal.

The fix was almost embarrassingly simple: raising shared_buffers to 2GB, enough to keep the working dataset fully resident in Postgres's own cache, cut the false-signal rate roughly 4x and also made query execution itself faster across the board. Any team building an RL environment around a real, stateful system with its own caching behavior — not just databases — should treat "how much does my own measurement noise pollute my reward signal" as a first-class engineering question, not an afterthought.

Honest limitations

• Narrow evaluation scope. All results are measured against a single 8.5GB dataset (IMDb) and read-only join queries — not representative of production OLTP workloads with concurrent writes, changing statistics, or larger data volumes, a caveat raised directly in Hacker News discussion of the writeup.
• Retraining cost isn't amortized in the headline numbers. The reported speedups don't include the ~95 hours of GPU time spent training the model in the total "cost per query" calculation — a fair critique raised in community discussion.
• Hints can become stale. If underlying data distributions shift meaningfully, previously-learned hints could regress performance rather than improve it — the model would need periodic retraining to stay valid, similar to how query plans generally degrade as data shapes change.
• Not a drop-in production tool. This is a research writeup demonstrating feasibility, not a packaged product — reproducing it requires building your own harness, dataset, and training pipeline.

What this means for what you build or pay

Teams with recurring, expensive analytic queries: this is a legible blueprint for a genuinely practical use of small-model RL specialization — train once against your own schema and query patterns, then run cheaply and repeatedly, rather than paying frontier-model inference costs on every query execution.

ML engineers exploring agentic RL generally: the reward-hacking and GRPO-mean-reinforcement failure modes documented here recur across almost any agentic RL setup, not just databases — worth reading in full even if query optimization itself isn't your domain.

Anyone deciding between "just call a frontier model" and "train a small specialist": Bansal's own conclusion is the right framing — validate feasibility with a frontier model first, then decide if the amortized cost of training and running a small model beats paying per-call frontier inference at your actual query volume.

Related on explainx.ai

• What is a database? How it works, beginner's guide
• Scalable oversight: RLHF, Constitutional AI, weak-to-strong generalization
• Xiaomi MiMo-V2.6: livestreaming a trillion-parameter RL run
• Castform Neon: RL post-training beats GPT-5.6 Sol at retrieval
• What are LLM parameters? Top 10 model sizes explained
• Top 10 open-weight models for a laptop

Original source: Rohan Bansal, "Training a 4B model to produce 81% faster query plans than Postgres", rohanbansal.com (September 2026)

Figures reflect the original writeup's reported benchmarks on the Join Order Benchmark as of September 16, 2026. Results are specific to the IMDb dataset and read-only analytic queries — see the honest limitations section before generalizing to other workloads.