Postgres is in a tough spot here. It would be reasonable to think it could simply count cardinalities and pick the plan that minimizes the number of rows passed through to successive joins.
But this would imply Postgres can count cardinalities during query planning. It can’t. In order to know this, it would need to actually run each join and count the resulting rows. This defeats the whole point of a fast query optimizer. A query optimizer does not aim to be exact in its cost minimization… it aims to be good enough across many types of queries.
Instead, Postgres uses statistics to estimate cardinalities. The planner queries the pg_statistic table, getting back common values for each column and their frequencies, and a histogram for the rest. Things get a bit more complicated when you tack on joins. Postgres doesn’t know how the rows in one table are distributed over the other. To get around this, it assumes that the frequency of a given value in the first table can simply be applied over the second table. This is the uniform distribution assumption I mentioned earlier.
Assuming a uniform distribution is fine as a heuristic, but when it fails, it fails hard. Looking back at an earlier join ordering (cn′⋈mc)≈100k, then ⋈ t′≈20k(cn’ \bowtie mc) \approx 100\text{k}, \text{ then } \bowtie\ t’ \approx 20\text{k}(cn′⋈mc)≈100k, then ⋈ t′≈20k, we filtered 2m movie_companies entries on the assumption that 5% of them were from Japanese companies. But what if the 5% of companies that are Japanese were actually responsible for 50% of the movies? The first join would produce 1m rows! The cost model says pick the first join ordering; in reality, the second one is actually better since it only sends 400k rows through to the second join.
One bad estimate in an early join can cascade through the rest of the join tree, corrupting all other estimates.
Postgres always picks the plan with the lowest cost, and we can’t change its cost model without modifying its source code, so how can we actually steer it to pick different plans that have higher costs?
pg_hint_plan is a beautifully simple third-party extension: just by adding structured “hints” as comments above SQL statements, you can nudge Postgres towards plans that use the instructions provided in the hint. For example:
/+The hint block. It is an ordinary SQL comment with a leading +, so Postgres ignores it and pg_hint_plan reads it. HashJoin(a b)Join a and b using a hash join. SeqScan(a)Read table a with a sequential scan rather than an index./EXPLAINEXPLAIN prints the plan Postgres would use instead of running the query. SELECT * FROM pgbench_branches b JOIN pgbench_accounts a ON b.bid = a.bid ORDER BY a.aid;
QUERY PLAN--------------------------------------------------------------------------------- SortThe root of the plan. Rows flow upward, so this runs last: it orders the joined rows by a.aid. (cost=31465.84..31715.84 rows=100000 width=197)Postgres’s estimates for this node: startup cost..total cost, estimated rows out, and average row width in bytes. Sort Key: a.aid -> Hash JoinThe join method the HashJoin(a b) hint asked for. (cost=1.02..4016.02 rows=100000 width=197) Hash Cond: (a.bid = b.bid)The join condition, taken from the ON clause. -> Seq Scan on pgbench_accounts aThe scan method the SeqScan(a) hint asked for. This is the probe side: each row looks up a match in the hash table. (cost=0.00..2640.00 rows=100000 width=97) -> HashThe build side. The small table is read first and loaded into an in-memory hash table keyed on bid. (cost=1.01..1.01 rows=1 width=100) -> Seq Scan on pgbench_branches b (cost=0.00..1.01 rows=1 width=100)(7 rows)
Example from pg_hint_plan’s
documentation.
The hint mandates usage of a HashJoin for joining pgbench_accounts and pgbench_branches, and doing a sequential scan of the pgbench_accounts table; the actual query plan follows suit nicely.
Given that we can influence Postgres to pick different—and potentially better—query plans using pg_hint_plan hints, the question we’re starting with is:
Can a language model learn to produce hints that result in better query plans?
What might make this a worthwhile problem to solve?
My first idea was to give the model the query and the exact same set of information Postgres’s planner has. This amounts to seeing if we could build a better cardinality estimator. I came to the conclusion this is not a worthwhile avenue to explore; we would be fighting decades of cardinality estimation research. Furthermore, the inference latency alone would far outweigh any learned usefulness compared to Postgres’s ultra-fast query optimizer.
The second idea—and what I believe is the correct formulation—lies in a specific database usage pattern: heavy analytic workloads. If queries are getting run thousands of times using sub-optimal default Postgres plans, efficiency gains are being left on the table. Instead, a model could be trained to find a better way to run a specific query. The training process might require execution of that query tens to hundreds of times upfront, but the amortized cost across all runs of the query would be drastically lower.
The goal isn’t to try and beat Postgres on the time/efficiency Pareto frontier for one-off queries, but we may be able to beat it on queries that run over and over again.
I decided to start with a small 4B model because it would be easiest to train/inference myself on the 2x RTX 3090 rig (affectionately named FLOPper) I have at home.
Around the time I started this project, the Qwen 3.8 family of models was released, unfortunately without a 4B variant. However, I came across a Qwen 3.8 4B distillation from a small lab in Germany called Empero and was intrigued. They used Qwen 3.8’s 2.4T model as a teacher model to distill learnings into Qwen 3.5 4B, producing empero-ai/Qwen3.8-4B-Distill. This distilled model is not outright better than its base 3.5 model; it performs better on MMLU tasks and slightly worse on GSM8K tasks. In other words, this distillation performs better when evaluated on breadth of general knowledge, and slightly worse on multi-step mathematical reasoning. As to which is better for our task, I do not know; I decided to stick with the distilled model either way.
With the model locked in, I built a lightweight agent harness, qo-agent, that would orchestrate hint production. It was given the following six tools:
To take advantage of structured outputs, the agent was instructed to produce PlanAction JSON objects. Calls to evaluate_candidate then compiled PlanAction objects into hints and prepended them to the original query.
Agent →A tool call made by the model. get_plan(“default”)Ask the harness for the plan Postgres picks on its own. This returns the planner’s estimates only; nothing is executed.Tool ←The harness’s reply, fed back into the model’s context. Default plan: (t ⋈ mc) ⋈ cn, hash joins, estimated rows per node Agent → evaluate_candidate({Propose one PlanAction. The harness validates it with a plain EXPLAIN, compiles it to hints, and, with execution feedback on, runs the query once after a warmup and reports the timing. “leading”: { “left”: { “left”: “cn”, “right”: “mc” }, “right”: “t” },The join tree: join cn with mc first, then join the result with t. This is the first tree from the figure earlier. “joins”: [{ “relations”: [“cn”, “mc”], “force”: “hash” }]Force a hash join for the cn ⋈ mc step. })Tool ← c1The id the harness assigns to this candidate. get_plan and finish refer to candidates by id. · valid · novel plan · 118 ms · 0.91× defaultExecution feedback: the candidate’s measured time and its preliminary ratio to the default’s median. Below 1 means slower than Postgres. · 2 attempts leftEvery submission, valid or not, consumes one of the rollout’s candidate attempts. Agent → evaluate_candidate({ “leading”: { “left”: { “left”: “t”, “right”: “mc” }, “right”: “cn” },Try the other tree: t with mc first, then cn. “joins”: [{ “relations”: [“t”, “mc”], “force”: “nestloop” }],A nested loop for the first join, so the filtered t rows probe mc one at a time. “scans”: [{ “relation”: “mc”, “force”: “index” }]Make the inner side of that loop an index scan on mc. })Tool ← c2 · valid · novel plan · 87 ms · 1.24× defaultFaster than the default: a ratio above 1 is a speedup. · 1 attempt left Agent → evaluate_candidate({ “leading”: { “left”: { “left”: “t”, “right”: “mc” }, “right”: “cn” }, “joins”: [{ “relations”: [“t”, “mc”], “force”: “hash” }]Same tree as c2, but a hash join for the first step, to see whether the nested loop was the win. })Tool ← c3 · valid · novel plan · 100 ms · 1.08× default · 0 attempts leftThe candidate budget is spent. From here the only tool the model is offered is finish. Agent → finish({ “selected_candidate_id”: “c2” })End the search by naming a candidate, including an earlier one. The harness then measures it against the default with the final paired protocol.Tool ← Finished · selected c2
A sample trajectory where the agent is permitted to submit up to three candidates.
Benchmarks
An agent is useless without something to benchmark its performance against. Fortunately for us, the hard work of creating these benchmarks was already done.
Leis et al. introduced the Join Order Benchmark (JOB) in How Good Are Query Optimizers, Really?. They used it to evaluate cardinality estimation and join-order optimization using our familiar IMDb dataset.
It consists of 113 queries spread across 33 query templates. Query templates differ via their relational skeleton. They reference different tables and connect them with different join predicates. You can think about them as a structural family of questions that can be answered. Queries derived from templates preserve the tables used and the join graph topology but change selection predicates.
SELECT MIN(t.title) AS movie_title FROM company_name AS cn JOIN movie_companies AS mc ON mc.company_id = cn.id JOIN title AS t ON t.id = mc.movie_id JOIN movie_keyword AS mk ON mk.movie_id = t.id JOIN keyword AS k ON k.id = mk.keyword_id WHERE cn.country_code = :country_code AND k.keyword = 'character-name-in-title';
Query template 2 — “What is the alphabetically first title of a movie associated with a company from country X and tagged with the keyword character-name-in-title?”
…and here are two real queries from JOB derived from this template:
SELECT MIN(t.title) AS movie_title FROM company_name AS cn, keyword AS k, movie_companies AS mc, movie_keyword AS mk, title AS t WHERE cn.country_code = '[de]' AND k.keyword = 'character-name-in-title' AND cn.id = mc.company_id AND mc.movie_id = t.id AND t.id = mk.movie_id AND mk.keyword_id = k.id AND mc.movie_id = mk.movie_id;Query 2a — “What is the alphabetically first such movie title associated with a German company?”
SELECT MIN(t.title) AS movie_title FROM company_name AS cn, keyword AS k, movie_companies AS mc, movie_keyword AS mk, title AS t WHERE cn.country_code = '[us]' AND k.keyword = 'character-name-in-title' AND cn.id = mc.company_id AND mc.movie_id = t.id AND t.id = mk.movie_id AND mk.keyword_id = k.id AND mc.movie_id = mk.movie_id;Query 2d — “What is the alphabetically first such movie title associated with a U.S. company?”
Another relevant benchmark is the Cardinality Estimation Benchmark (CEB), introduced in Flow-loss: Learning Cardinality Estimates That Matter. It uses the same IMDb database and is a much larger benchmark consisting of ~13.6k synthetically generated queries organized across 16 query templates CEB’s definition of a template is looser than JOB’s. Two CEB templates can share the same join graph, differing only in their selectivity predicates. In JOB, every template’s join graph is unique. .
Due to its size, CEB was a good fit for training the model. JOB would be used to validate the model’s performance.
You might be wondering if it makes sense to both train and test on IMDb. If it works well, hasn’t the model just learned this specific database well?
I would argue this is precisely the point. We want our model to learn IMDb well. Given our problem formulation, if this agent is continually getting used for a company’s analytic workloads across its specific databases, we need not generalize to all databases.
The real issue is making sure we’re not overfitting to JOB query templates during training over CEB. The model should learn IMDb in a way where given any query, even for structural query families it hasn’t seen before, it’s still capable of producing a good plan. In practice, this means we need to prune CEB queries that have the same shape as any of the JOB queries.
Let’s define a query’s “topology” as its structural join-graph (de-aliased table names as nodes and joins as edges). The join graph excludes all selectivity predicates; we’re only interested in joins here.
CEB queries sharing a topology with a JOB query would be removed from the training set. I wrote a small script to convert all JOB and CEB queries to their topologies and checked if there was any overlap. There wasn’t, so no filtering was required.
JOB: 113 queries, 33 templates, 33 topologies
CEB: 13,646 queries, 16 templates, 12 topologies
JOB and CEB templates displayed as an identicon of their topologies.
Before getting into benchmarking the agent and doing training runs, we have to talk about how Postgres was actually run, because it directly impacts the training process.
If we run the exact same query on Postgres 20 times in a row, it won’t take the same amount of time each run. In day-to-day work, this isn’t a big deal. But the whole thesis, and the training process itself, relies on measuring whether one way of running a query is faster than the Postgres default. This means we need to do everything in our power to de-noise Postgres.
First, I needed to understand just how noisy Postgres query executions are.
I started by building a “calibration” capability into my experimentation workflow. The calibration process was simple: run NNN Docker containers built from a Postgres image, each given a fixed slice of CPU cores and RAM to use. I set N=4N = 4N=4 to begin; anything lower might make future training far too slow, and anything higher might lead to more CPU contention, which means more noise. Each container was given 4 cores to use and capped at 8 GB of memory.
On startup, each container initialized Postgres with identical settings and loaded the IMDb data. Calibration then opened a thread pool of size four and pushed all 113 queries onto a shared queue. Whenever a container finished measuring a query, it pulled the next one off the queue.
The actual measurement process had two phases:
So what does it mean to warm a query up? We need to bust out some OS fundamentals to understand.
Whenever Postgres executes a query, it asks the operating system (in our case, Linux) for pages of data. Linux first checks its own filesystem cache, the page cache. If the pages are present, Linux sends them over; else it reads them from disk, stores them in its cache and then sends them over. Postgres, in turn, keeps received pages in its own shared_buffers cache for easy reuse. When shared_buffers begins to overflow, Postgres evicts pages. If it needs those pages again, it must ask Linux once more.
Every time there’s a cache hit in shared_buffers for a page, Postgres increments a counter called “shared hit blocks” (SHBs). If it has to ask Linux, it increments “shared read blocks” (SRBs).
Postgres conveniently reports both counters if we run EXPLAIN with the BUFFERS option. For example, running EXPLAIN (ANALYZE, TIMING OFF, BUFFERS, FORMAT JSON) outputs something like:
{ "Plan": { "Node Type": "Aggregate", "Shared Hit Blocks": 1800786, "Shared Read Blocks": 52990, ... }, "Execution Time": 189.2, ..., }
These counters give us some notion of the “warmness” of a query. After each warmup run, we compared its hit and read counts to the previous run’s. If both were within 2% of each other (and the plan hadn’t changed), we called the query warm and started measuring. A query needed at least two warmups to have something to compare, and was cut off at five regardless. The idea was that if the counters stopped moving, the data could be considered settled and cache churn would be minimized during the 20 measurements.
Query A Query B Shared hit blocks 0 Shared read blocks 0
I set shared_buffers to a conservative 128 MB and ran the first calibration:
Still reading from Linux after warmup Fully resident in shared_buffers
Half the queries were declared warm after only two runs. Not bad… at least until I dug deeper. The SRB counts weren’t dropping to zero; rather, they were hovering steady at some large number. With only 128 MB of shared_buffers against an 8.5 GB database, Postgres was consistently missing its own cache on every execution and asking Linux for more pages. “Stable” did not mean “resident.”
Linux’s page cache is fast, so this isn’t the end of the world. Unfortunately, a new problem emerged when I actually looked at the 20 measurements taken for various queries. Let’s look at one query in particular, job-13b:
14 of the 20 landed between 186 and 204 ms. The other 6 landed between 227 and 253 ms, somewhere between 14% and 26% slower. The query wasn’t even uniformly noisy, it just had two different speeds at different times, and a third of the time it ran at the slower speed.
I initially wanted to quantify noise using the coefficient of variation:
The CV tells us the “wobble” of a measurement. If a query takes 100 ms and has a CV of 5%, we could say it wobbles by about 5 ms. For job-13b, the CV was 10.3%. It wasn’t great. CV is also not a great measurement to use here. Because it’s built on the mean, it’s easily influenced by a few outlier runs.
We don’t actually care as much about how spread out the 20 runs are. We do care about how often this causes our measurement criteria during training runs to get fooled.
Bear with me here as I skip ahead a little bit in order to provide more color on what exactly we needed to measure.
To de-noise during actual agent runs, I couldn’t just run the agent’s proposed plan a single time. Instead, I ran three interleaved (candidate, default) pairs sequentially. Three was picked somewhat arbitrarily to provide some measure of variability while being small enough to prevent agent evaluation runs from spending most of their time in Postgres. Once the three candidate/default execution time tuples were obtained, the medians of both the three candidates and the three defaults were taken and expressed as a ratio of each other to determine the final speedup or slowdown. If the two medians differed by less than an arbitrarily declared 5%, it was a tie. Outside of that tie zone, a candidate could be declared as a speedup or a slowdown.
Now let’s go back to our earlier job-13b example. We had 14 executions in one clump, and 6 in another slower clump. The median of three strategy sounds good until you realize that if, in theory, at least two of the three measurements landed in that “slower” clump, the median would bias towards the less frequent slower clump.
Imagine a candidate plan that executes identically to the default. No real difference exists, so the correct reward is zero. Draw three timings for the “candidate” and three for the “default” out of the 20 we observed. There are (203)=1,140\binom{20}{3} = 1{,}140(320)=1,140 ways to draw three from 20; for job-13b, 230 of them contain at least two slow runs, so one side’s median lands in the slow clump ~20% of the time.
That’s a totally phantom 14-26% speedup or slowdown that we would show to our model as signal ~20% of the time. Dangerous!
job-13b 128 MB shared_buffers · a no-op candidate (i.e. one that is identical to the default)
0 rounds · ties 0 · phantom wins 0 · phantom losses 0 · fooled 0%
So we can’t just rely on CV as the golden number to minimize, as two queries with the exact same CV can fool the measurement reward at different rates depending on whether the spreads are a uniform blur or two clumps sitting more than 5% apart. The actual number to minimize is this fooling rate itself.
I wrote a small script to compute the fooling rate directly from raw calibration data. It worked by sliding a window of six sequential runs across the 20. For each window, we took interleaved pairs of size two to represent an interleaved (candidate, default) pair. A window of size six gives us pairings like: (t1, t2), (t3, t4), (t5, t6). In any given pair, tnt_ntn and tn+1t_{n+1}tn+1 can alternate roles of being the candidate query, or the default query. That means for each pair, there are two possibilities, and therefore for each window of three tuples, there are 2×2×2=82 \times 2 \times 2 = 82×2×2=8 possibilities. 20 measurements means we’ll slide this window 15 times, so we have 15×8=12015 \times 8 = 12015×8=120 total possibilities Each possibility is a binary value indicating whether or not that specific, simulated formulation of candidate/default pairs resulted in a ratio of medians between the two greater than the 5% tie-zone. for a given query.
We derive two metrics from these raw numbers. First, we calculate the no-op error rate for a given query as the ratio of the 120 simulated possibilities that do differ by more than 5% against the number that don’t. We sum these percentages up across all 113 JOB queries and then divide by 113. This number, which we’ll call the “mean no-op error rate,” gives us the percentage likelihood that the reward may get fooled for any JOB query when doing our three paired measurements strategy. Second, we sort the no-op error rates for all 113 queries, lowest to highest. The number that is 90% of the way to the end of this sorted list is reported as the “p90 query,” and gives us a measure of the fooling rate for the worst-offending queries.
At 128 MB for shared_buffers and four concurrent containers, the “fool rate” script produced the following mean no-op error rates and p90 query numbers I ran the calibration twice per config to provide a sense of how much two runs may disagree with each other. :
The numbers aren’t good. One in twenty no-op plans get rewarded, and one in ~10 queries gets fooled more than 13% of the time.
I focused on two memory-related settings Postgres exposes:
Surprisingly, work_mem had no effect on noise at all, and shared_buffers carried all of the weight!
With 2 GB of shared_buffers, the median query ended warmup with its SRB counter at exactly zero: its working set was fully resident in Postgres’s own cache. The no-op error rate dropped by roughly 4x, and the 90th percentile query went from being fooled 13%–20% of the time to almost never. Our two-clump query, job-13b, went from a CV of 10.3% to 0.9%, with all 20 runs landing within 7 ms of each other.
One neat benefit emerged that I wasn’t initially chasing: the default plans themselves got faster. The summed runtime of all 113 JOB queries fell from 95 seconds to 60 seconds, just from cache residency. In other words, actually taking our measurements for both candidates and defaults would now be significantly faster, meaning the training process would take less time.
I locked in 2 GB shared_buffers and 4 MB work_mem for the rest of the project.
I used two metrics for benchmarking agent performance.
The geometric mean speedup gives all queries equal weight. For example, in a two-query sample, if query 1 runs 2x faster than its baseline, and query 2 runs 0.5x faster than its baseline, then Sgeo=1.00xS_{geo} = 1.00\text{x}Sgeo=1.00x. It doesn’t matter if query 1’s baseline took 5 minutes and our candidate took 2.5 minutes, but query 2 only regressed from 25s to 50s, as they are equally weighted.
Total workload speedup treats the entire query set as one batch. We simply add all the baseline times and divide by the sum of the candidate times. In our above example, Sworkload=1.4xS_{workload} = 1.4\text{x}Sworkload=1.4x.
Both metrics tell different stories. The total workload speedup is a measure of practicality. A data analyst building out a suite of analytics queries wants to decrease the overall runtime across the batch. But from a model training standpoint, the total workload speedup could be entirely influenced by a single query plan the agent chanced upon; the rest of the batch could be degenerate. This implies the model hasn’t actually learned anything interesting; it just got lucky. Because the geometric mean speedup cares not for absolutes, it gives us a measure of actual learning across the batch: values above 1x imply that the average query is executing faster.
Before running the untrained 4B model through the qo-agent harness, I wanted to validate this problem was actually solveable by today’s frontier models. If a model like GPT-6 Astra or Qwen 3.8 2.4T couldn’t improve upon the default Postgres query plan, I couldn’t really expect the 4B model to either.
I took a small sample of 10 JOB queries and benchmarked them on both Astra and Qwen 3.8 2.4T running through the qo-agent harness:
Evaluations of Astra and Qwen 3.8 2.4T run on the same slice of 10 JOB queries. The frontier models were benchmarked at different candidate numbers (i.e. how many candidates they were allowed to generate during a complete trajectory; either a single candidate or 5) and for Astra, whether reasoning summaries I was a little surprised to see Astra performance worsen with reasoning summaries on compared to the 5-candidate evaluation done right before it, but these evaluations were only run a single time on a small 10-query slice of JOB, so I chalked up the worse results to random variance. were enabled or not. Astra was inferenced through OpenAI’s API, and Qwen 3.8 2.4T through Modal via OpenRouter.
Given the difference between the single-candidate scores and the 5-candidate scores, the agent was clearly capable of doing in-context learning across sequential executions of its candidates. This gave me the confidence to stick with an agentic multi-turn approach rather than try and train the 4B model to get really good at one-shotting a plan.
During a run of the agent, each candidate was warmed once and then measured once. After exhausting the candidate attempts budget, the model was only presented with a single tool to call, finish, and the model was told to select the best scoring candidate (or keep the default plan). After the candidate was selected, three interleaved (candidate, default) pairs were run and passed through a clipper:
The clipper constrained the result of the division between the two medians to be between [0.1,10][0.1, 10][0.1,10]. These clipper values were picked somewhat arbitrarily; I found they prevented the geometric mean speedup from getting overly influenced by an extreme speedup or an extreme regression.
Conclusion: frontier intelligence is capable of agentically doing query optimization.
We’re now ready to evaluate the untrained 4B model on JOB and see how it does!
The same 5-candidate plan budget per agent trajectory configuration was employed. The results were dismal:
Only the last three rows contribute to the score, leaving 15 valid trajectories out of 113. Nine of those 15 result in a score of 1.00x by construction (the plan was identical A candidate plan was determined to be equal to the default plan if their EXPLAIN outputs with cost and row estimates stripped were equivalent. to the default, or the model chose to keep the default). That left just six candidate plans that were:
Five of the six plans resulted in speedups of 1.02x–1.30x, and one of them landed at 0.05x.
Not only was the model terrible at this task, it couldn’t even grok the harness wrapped around it either.
I first needed to get the 4B model to speak the “language” of the qo-agent harness. I would make it good at query optimization after.
We can use supervised fine-tuning (SFT) to do this. Specifically, we can do off-policy distillation.
Off-policy distillation is a training method by which a student model (sometimes referred to as a policy) learns to imitate outputs produced by a teacher model. It’s called “off-policy” because the training data is not generated by the student model/policy itself. It’s remarkably simple. A complete teacher trajectory (sometimes referred to as a demonstration) is shown to the student model. For every token in the trajectory, the probability the student model gave to generating that token results in a per-token loss. Averaging these per-token losses leads to a demonstration-level loss. Standard backpropagation via chain rule then lets you compute gradients for all trainable parameters in the student model, and the configured optimizer can nudge parameter values in a way where loss gets minimized in a single pass.
What’s super neat about off-policy distillation is that we don’t need a lot of data for it to work well. Because each demonstration provides us thousands to tens of thousands of token predictions, our student model’s weights adapt quickly.
How do we actually get these demonstrations though? We could write them all by hand, but that would take far too long. One step up would be writing a tool to randomly generate valid-looking trajectories Fun fact: I tried this initially. It actually works decently and was able to teach the 4B model the harness. However, it caused a bunch of other issues, mostly around making the model less likely to generate novel candidate plans. . But these ignore that we have the best teacher of all already available: smarter, larger models.
The strategy is simple: generate a bunch of trajectories by running a smart model through the qo-agent harness, and distill those trajectories into our student 4B model.
Imagine we decide to use GPT-6 Astra as our teacher model. It produces trajectories in OpenAI’s Responses API format. Our Qwen model doesn’t understand this format; we need to transpile the human-friendly Responses API JSON format into a model-friendly token format. This is where the concept of rendering comes in. Rendering libraries can take in a trajectory’s text and convert it to a raw sequence of tokens a specific model actually understands.
Another consideration with agent trajectories is determining which tokens our model should actually be predicting. The model never produces certain tokens in a trajectory, like the system prompt, any user prompts or the results of a tool call after a harness executes it. The model should still see these tokens when predicting the next token though; they’re still part of the context, but we should only compute losses for tokens the model is responsible for predicting. We can employ a strategy called loss-masking here. A loss-masking library lets us label the parts of a teacher trajectory that are context-only, versus the parts our student model is responsible for predicting.
Finally, when fine-tuning over agent trajectories, we generally don’t include the entire trajectory as a single trainable unit. Instead, the trajectory is broken up into a set of (context, reply) pairs. The context in these pairs is additive and includes previous model replies.
Tokens the model is scored on Context only, no loss
Our 4B model has 4.66 billion tunable parameters. If we wanted to update all of these parameters in a single pass during SFT, we would need a GPU with at least 64 GB of VRAM. I wanted to test my hypotheses first on my consumer-grade RTX 3090s, each of which carries 24 GB of VRAM, so I needed something more parameter-efficient here.
Low-rank adapters, or LoRAs, are the canonical way to do this. At a high-level, they work by freezing the model’s weights as they are, instead letting you train a much smaller pair of matrices that when multiplied together, result in an adjustment to selected weights in your original model.
4.66 billion weights come in at 9.32 GB in bf16. The tiny LoRA I actually ended up training was only 42.5 MB and contained only 21.2 million trainable parameters.
We previously learned that frontier models could operate well in qo-agent; the next decision was picking between GPT-6 Astra or Qwen 3.8 2.4T.
Let’s look back at the results from evaluating both models on 10 JOB queries, this time zooming into context length in tokens:
The same five-candidate evaluations over the 10 JOB queries mentioned earlier, this time measured by how much context each trajectory consumed.
It seems like Astra wins on all fronts. However, Astra’s biggest drawback is that when run through the API, reasoning tokens aren’t supplied. Instead, the API gives you a short summary in place of reasoning tokens. On the other hand, Qwen 3.8 2.4T is an open-weights model and is happy to provide all of its reasoning tokens.
I was wary about training off of trajectories containing only reasoning summaries. The How to Steal Reasoning Without Reasoning Traces paper talked about this exact thing: training off reasoning summaries resulted in the student model’s performance decreasing. To combat this, the authors devised a new method: trace inversion. Trace inversion calls for synthetically expanding a reasoning summary into what the raw reasoning tokens may have looked like. Although they’re not exactly the ones Astra actually produced during inference, the longer reasoning blocks led to improved performance when transferred over to a smaller model. This provided some level of comfort; if I picked Astra and performance suffered, I could experiment with trace inversion.
The other consideration was context lengths. Given I was only doing SFT off a single RTX 3090 to start, I needed any given trajectory to not exceed ~50k tokens in sequence length. If it did, the training process might OOM given the 3090’s limited 24 GB of VRAM. All Astra trajectories fit under that budget, but some of the Qwen 3.8 2.4T ones didn’t.
I decided to try out SFT on the Astra traces. If performance suffered, I could explore trace inversion; if that didn’t work, I could rent beefier GPUs for training, bump our global context limit, and use Qwen 3.8 2.4T traces instead.
I began by generating 120 Astra trajectories over a random slice of queries from CEB, with reasoning summaries enabled. 100 trajectories would be used for training, and 20 would be used as a held-out validation set. All trajectories were rendered via Prime Intellect’s renderers library into Qwen format, loss-masked appropriately and unrolled and packed into usable training demonstrations. I then used Prime Intellect’s prime-rl library to run SFT, training the smaller ~21 million parameter LoRA. 100 Astra trajectories became 382 training rows after unrolling and packing. I allowed training to run for just a single epoch, meaning every example was trained on exactly once, and used a batch size of one (each demonstration
The one-epoch adapter against the untrained model on all 113 JOB queries. A query counts as scored when its trajectory ended with a measured candidate, a duplicate of the default plan, or the default itself. Wins and regressions are scored queries more than 5% faster or slower than the default.
Promising! The adapter learned the harness.
I could now either generate more fresh trajectories, or do more epochs over the dataset we already had. Given the latter is cheaper, I decided on more epochs.
It was around this point that I got impatient and wanted the training process to run even faster, so I rented a 2x H100 node on Lambda.
Now training on an H100 Switching to the H100 meant we had 80 GB VRAM at our disposal instead of 24 GB. We could have bumped the context token limit greatly, but I decided to barrel through with the roughly ~50k token limit we set. instead of a single RTX 3090, I did two more training runs Training runs were a lot faster on the H100. The first epoch took four hours on the RTX 3090; the second took only 45 minutes on the H100. over the existing LoRA; a second epoch and then a third:
Two and three epochs over the same 100 Astra trajectories, evaluated on JOB. Before epochs two and three, the qo-agent harness was upgraded to allow the model to keep the default plan after a search, which is why many more queries are scored. The two- and three-epoch rows were evaluated under identical settings.
Two epochs improved our results, but three epochs regressed them! This was especially interesting because validation loss didn’t budge at all during the second epoch:
Validation loss on the 20 held-out trajectories
0.485 0.305 0.306 0.322
0 382 764 1,146
epoch 1 epoch 2 epoch 3 Optimizer updates, one per packed training row
A very good lesson that a flat validation loss doesn’t necessarily mean the model has stopped learning useful behavior.
I still felt we had more to learn from SFT though before proceeding with RL. I did zero filtering on the training trajectory dataset, and hadn’t carefully audited if I was missing any capabilities. It turns out I was, mostly around the model’s ability to construct valid Leading trees.
I generated another 320 Astra trajectories; 300 for training, and 20 for validation. I filtered out just six training trajectories where Astra opted to keep the default plan without even trying a single candidate. I did two more epochs in two separate runs:
Continuing the two-epoch adapter on the 300 new trajectories, evaluated on JOB.
Our 4B model not only learned the harness; it was now genuinely making good calls on various JOB queries! Fortunately, training on reasoning summaries didn’t harm performance.
The model now spoke the “language” of the qo-agent harness, and we got some free performance gains out of SFT too. It was time to make it very good at query optimization.
Agentic RL differs from SFT in that we actually run the current policy over training queries inside the agent harness NNN times. Each run (also referred to as a rollout) results in a final output that’s scored against some verifiable criteria. Lastly, each rollout’s score is then weighted relative to the other same-query rollouts. A positive “advantage” is reinforced by making the model’s weights more likely to produce that trajectory in future runs, and a negative advantage is penalized; the weights are The speedup was calculated as the median of three default plan measurements divided by the median of three candidate plan measurements.
For every evaluated candidate that was invalid, we subtracted 0.1 from the natural log of the speedup ratio. We subtracted a further 0.05 if the rollout resulted in a plan that shared the same fingerprint as the default Postgres plan. Finally, if the trajectory ended with no valid candidate at all, a flat 3 was subtracted from the reward in lieu of any of the 0.1 or 0.05 subtractions.
The first few RL runs I did using this reward resulted in a model that was terrified of producing invalid plans due to the extremely harsh -3 condition. The model played it safe instead, returning Postgres’s default plan over and over again, accepting the smaller 0.05 reward hits.
GRPO exacerbated this issue. The plain GRPO algorithm converts multiple rollout rewards into relative “advantages”:
GRPO as prime-rl implements it: simply
subtract the group’s mean reward from each rollout’s reward. The GRPO paper also
divides by the group’s standard deviation.
Let’s say we perform four rollouts for a given query resulting in the following plans, execution speeds and rewards: