Accelerate Deep Research Workflows with Amazon Bedrock AgentCore, Part II: Customising the Model
Part II adds model customisation to AgentCore Deep Research. You can distil a frontier teacher's research behaviour into a small open model and then sharpen it with reinforcement learning, using the agent you already have as both the data source and the training environment.
GitHub repo: https://github.com/aws-samples/sample-agentcore-deep-research
Full results and reproduction steps: SFT distillation and RL fine-tuning
Part I showed how to build an agent that scaffolds a report, researches across 10+ connectors behind AgentCore Gateway, writes it with inline citations, and verifies its own coverage. It runs on a frontier model through Amazon Bedrock.
This post adds a pipeline that trains a small open model to run that same workflow. Everything described here is in the repository and runs end to end, including the data generation, the SFT job, the RL job, the serving stack, and the evaluation harness.
1. Why customise the model
A deep research episode is not one model call. In our agent a single question produces 20 to 30 turns of tool use, and the conversation grows monotonically, because every search result stays in context so the model can cross-reference it while writing. Measured across 1,833 real trajectories, the median episode is 45 messages and roughly 26,000 tokens, two thirds of which are tool observations. That shape is what makes customisation worthwhile, for four reasons.
Cost per episode. You are not paying for one frontier completion. You are paying for 20 to 30 of them over a context that keeps growing. Serving a 9B model yourself on a single GPU changes the unit economics of the whole workflow rather than trimming a percentage off it.
Fewer tokens for the same behaviour. Our production agent is steered by a long system prompt that spells out the five-step workflow, the citation format, and which connector suits which kind of question. A model that has learned the workflow from trajectories needs less of that instruction repeated on every turn, because the behaviour sits in the weights instead of the prompt.
Latency. Turns are sequential, since the agent must see one search result before deciding what to search next. Per-turn latency therefore multiplies by 20 to 30, so it matters more here than in a single-shot application.
Your domain and your sources. This is the part a general-purpose model cannot give you. If your research questions concern your own filings, your own knowledge base, or a speciality corpus behind your own Gateway tools, then the behaviour worth learning is how to research your sources: which connector answers which kind of question, what your documents look like, and how your analysts expect a report to be structured. You generate that training data by running the agent you already have against your own tools.
The asset ships the pipeline, not a model. The repository contains the machinery to train a student model on your questions and your connectors, with the same agent code serving production, data generation, and RL rollouts.
2. How it works
There are two stages. Both run on SageMaker and both use the deployed AgentCore agent as their environment. Stage 1 teaches the behaviour and stage 2 optimises it against a score.
Stage 1: distil trajectories, not reports
sft_generate_data.py runs the production agent against a set of research questions and captures the whole trajectory from the response stream: every tool call and every tool result, in order. No server-side capture hook is needed, because Strands emits one message event per completed message and tool results arrive as toolResult blocks on those events.Training on
question → final_report pairs would instead produce a model that writes report-shaped prose with no research behaviour behind it, which invites fabrication. Every published approach in this area ( DR-Venus , DeepSearch-World , DeepRubric , R²-Searcher ) trains on full multi-step trajectories, with tool observations masked out of the loss. Without that mask the model is trained to generate search results, and it will duly invent them at inference time.Training uses TRL's
SFTTrainer with LoRA under FSDP, at max_seq_length 32768. A trajectory is roughly 26K tokens, which is the single biggest driver of the hardware requirement.Stage 2: reinforcement learning against a rubric
Stage 2 runs GRPO on the AgentCore RL Toolkit with the verl backend and its FSDP engine. The arrangement that makes this practical is that AgentCore Runtime serves the rollouts. Each training step sends prompts to the same agent container that runs in production, which calls the same Gateway tools, so the policy is exercised in the environment it will be deployed into rather than a simulation of it.
Inside the training job,
rllm-model-gateway sits between the agent and vLLM and captures token IDs and logprobs as the rollout happens, so GRPO receives exactly the tokens the policy emitted. Using verl with FSDP means the model architecture is read from config.json , so there is no per-model training script to maintain.Each report is scored from 0 to 1 on three signals, defined once in
research_rubric.py and shared by the RL reward and the evaluation harness:Signal | Weight | How |
|---|---|---|
Rubric quality | 80% | LLM judge over 6 criteria: coverage, sources, synthesis, depth, actionability, grounding |
Citation validity | 10% | Distinct well-formed URLs, rejected unless they appear in what the tools actually returned |
Format compliance | 10% | Every expected section carries real body text, not just a heading |
The two objective signals matter more than their weight suggests. In one run the judge scored a 13,000-character report containing zero URLs as "sources: 9/10", while the citation component correctly returned 0.000. Judges are unreliable at verifying provenance, and deterministic checks are good at it.
Three key learnings
Check that observation masking is actually on. TRL's
assistant_only_loss builds its mask from {% generation %} markers in the chat template, and not every published template has them: Gemma 4 ships them, Qwen3.5 does not. For recognised model families TRL handles this for you by swapping in a marked-up training template, so there is nothing to configure. The gap is an unrecognised family, where the flag does not error and simply trains on the whole sequence, observations included, which is the outcome it was meant to prevent. Observations are around half the tokens, so the loss curve looks plausible either way. The pipeline therefore asserts that markers come from either the model's own template or TRL's, and refuses to start otherwise. The same detail applies in reverse at serving time, since {% generation %} is training-only, so the merged checkpoint is written with the published template restored.Conversation history must not be truncated. Strands defaults to
SlidingWindowConversationManager(window_size=40) unless you say otherwise. Our episodes are p50=45 messages and reach 75, so 86% were being truncated mid-episode, sometimes splitting a tool call from its result. This cost every model quality, but it cost the small model far more than the frontier one, so it did not simply add noise. It changed which conclusion the numbers supported. Set conversation_manager explicitly.The reward must be as hard to game as the metric. Our citation component rejects URLs the tools never returned, and the judge is shown the retrieved material so it can check whether a claim is supported rather than merely attributed. Both require the tool observations to be threaded from the rollout into the scorer, and that wiring is easy to leave half-connected: the parameters existed and were plumbed through the scoring functions, but the call site passed neither, so the RL reward gave full credit for plausible invented URLs and the judge assessed grounding with no sources in front of it.
3. A demonstrative example
To show the pipeline works end to end we ran it once on general-domain research questions with Qwen3.5-9B as the student. This is an illustration of the machinery rather than a benchmark result, and the shape of the curve is more informative than the individual figures.
Every model runs through an identical harness: same agent code, same tools, same prompt, same judge, greedy decoding, on 98 held-out questions. The model is the only variable.
Two choices in that harness are compute optimisations worth naming, because they shape what the numbers mean. The judge is the same mid-size model used for the RL reward, chosen because it is fast and cheap enough to run inside a training loop rather than because it is the strongest available evaluator. The judge is shown the most recent 24,000 characters of retrieved source material rather than all of it, which caps judge input cost but means a claim sourced from an early search can be marked ungrounded because its supporting snippet is no longer in view.
Training also caps each tool observation at 1,200 characters to keep a 20-turn trajectory inside the context window, while evaluation passes observations at full length. That asymmetry works against the fine-tuned models, which learned to research against the shorter observations, so the distillation gain below is if anything understated.
Stage | Score | Citation validity | Training compute |
|---|---|---|---|
Qwen3.5-9B, untouched | 0.538 | 0.299 | n/a |
+ trajectory SFT | 0.744 | 0.953 | 138 GPU-h |
+ GRPO, 50 steps | 0.772 | 0.906 | 168 GPU-h |
+ GRPO, 100 steps | 0.770 | 0.929 | 198 GPU-h |
+ GRPO, 150 steps | 0.782 | 0.919 | 228 GPU-h |
+ GRPO, 200 steps | 0.677 | 0.946 | 258 GPU-h |
Teacher (frontier) | 0.743 | 0.966 | n/a |
n=98 per model, paired comparisons. SFT vs base: +0.208, 95% CI [+0.174, +0.242], p=1.2e-20. GRPO 150 vs SFT: +0.038, 95% CI [+0.016, +0.060], p=0.001.
Distillation does the work. The base model gains +0.208 from trajectory SFT. That is the result we would stand behind, because it is large and it survives every methodological correction we have applied, including re-scoring with a different and stricter judge.
The mechanism shows up in the citation column. The move from 0.299 to 0.953 needs no judge at all, since it is a deterministic check against URLs the tools actually returned. The base model emits unresolvable bare domains, while the distilled model learned the teacher's habit of citing pages it genuinely retrieved. Per-criterion judge scores move the same way, with "sources" rising from 3.69 to 7.32 out of 10.
The difference is visible in a single question. Asked about recent developments in the search for extraterrestrial life, the base model produces a fluent, confident report with zero citations, asserting specifics a reader has no way to check:
Base model, citation score 0.000 — “JWST has confirmed its first exoplanet, LHS 475 b, a rocky world nearly Earth-sized orbiting in just 10 days, and discovered Beta Pictoris d using a novel chemical fingerprinting technique.”
The distilled and RL-tuned model answers the same question with 17 citations, every one of them to a URL its tools actually returned, dates its claims, and says plainly where the evidence stops:
After SFT and GRPO, citation score 1.000 — “As of September 2026 ... No definitive biosignature has been detected, but the field has matured significantly, with the Virtual Planetary Laboratory joining the SETI Institute in July 2026 and the SETI@Home project completing analysis of 12 billion detections.”
Same harness, same tools, same question. The scores were 0.247 and 0.867. What changed is not writing ability, it is whether the model goes and finds out.
RL adds a further gain, then over-optimises. GRPO peaks at 0.782 by 150 steps, +0.038 over SFT (95% CI [+0.016, +0.060], p=0.001), against the teacher's 0.743. By step 200 the score falls to 0.677, which is 0.104 below step 150 (p<0.0001) and below SFT. The in-training reward was still climbing at that point, around 0.84 and rising, so the reward signal gave no warning at all. The failure is not uniform degradation: median report length barely moves, but the number of reports scoring below 0.5 on the rubric jumps from 4 to 20. A few episodes collapse rather than every episode getting slightly worse. Checkpoint often and evaluate each one; a rising training reward is not evidence that the policy is still improving.
The scale here is also far below what the published work uses. DR Tulu-8B , a comparable open 8B deep research model, reports 136 GPU-hours for its SFT stage and 9,700 GPU-hours for RL, covering 1,900 optimiser steps over 25 days on 16 H100s. Our SFT stage is essentially the same size at 138 GPU-hours, while our RL stage is 194 GPU-hours on 4x L40S. Measured in rollouts rather than hours, they train on roughly 486,000 episodes against our 1,200.
4. Try it on your own domain
The pipeline is in the repository. Point it at your questions and your Gateway tools.
Plain text
# 1. Generate research questions (seeded, reproducible)uv run test-scripts/sft_generate_questions.py --count 1000 --eval-count 98 --seed 7# 2. Collect teacher trajectories from your deployed agent, against your toolsuv run test-scripts/sft_generate_data.py \--questions test-scripts/results/rl_train_data_1k.jsonl \--max-concurrent 45 --observation-chars 1200# 3. Distil into a small open model (LoRA, observations masked)uv run test-scripts/sft_train.py --hf-model-id Qwen/Qwen3.5-9B \--lora-rank 32 --epochs 2 --max-seq-length 32768 \--instance-type ml.g6e.12xlarge# 4. Serve it, and point an agent at ituv run test-scripts/deploy_model.py --job-name <training-job> --endpoint-name dr-sftuv run test-scripts/deploy_finetuned_agent.py --endpoint-name dr-sft# 5. Optionally sharpen with RL against the rubricuv run test-scripts/rl_train.py --agent-arn <RLAgentRuntimeArn> \--sft-job-name <completed-sft-job> --instance-type ml.g6e.12xlarge# 6. Evaluate on your held-out questions, same harness for every modeluv run test-scripts/eval-agent.py --benchmark rubric --max-questions 98 --tag mine
One practical note. The rubric in
research_rubric.py is the piece you should expect to change, because ours rewards coverage, grounding and citation discipline for general research, and yours should encode what a good report means in your domain.The full walkthrough, including the evaluation harness and the reward definition, is in the repository README .
Conclusion
In this post, we showed you how to customise the model behind AgentCore Deep Research. The pipeline distils full agent trajectories from a frontier teacher into a small open model, then optionally sharpens it with GRPO reinforcement learning against a rubric reward, using AgentCore Runtime to serve the rollouts so the policy trains against the same tools it will use in production. Because the reward and the evaluation harness share one definition, you can change what a good report means for your domain in a single file and have both training and measurement follow.
To get started, clone the repository , generate trajectories from your own Gateway tools, and train a student model on your questions. To learn more about Amazon Bedrock AgentCore, see the AgentCore documentation , and for the RL backend see the AgentCore RL Toolkit .
Authors
Nikita Kozodoi, PhD is a Senior Applied Scientist at the AWS Generative AI Innovation Center, where he builds AI and agentic solutions for enterprise customers across industries, including Nasdaq, Ryanair and adidas. He holds a PhD in machine learning with publications at NeurIPS, KDD and ECML workshops, and has presentations at over 25 tech conferences.
Aiham Taleb, PhD, is a Senior Applied Scientist at the Generative AI Innovation Center, working directly with AWS enterprise customers to leverage Gen AI across several high-impact use cases. Aiham has a PhD in unsupervised representation learning, and has industry experience that spans across various machine learning applications, including computer vision, natural language processing, and medical imaging.
Jack Butler is an Applied Scientist at AWS, leading innovative projects at the AWS Generative AI Innovation Centre with a background applying large language models and applied AI research across a wide variety of enterprise and startup customers.
Liza Zinovyeva, PhD is an Applied Scientist at AWS Generative AI Innovation Center and is based in Berlin. She helps customers across different industries to integrate Generative AI into their existing applications and workflows. She is passionate about AI/ML, finance and software security topics. In her spare time, she enjoys spending time with her family, sports, learning new technologies, and table quizzes.
Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article