Context Optimization RL: a control system, a dataset, and a recipe

A frozen-model harness control system: an outer policy (contextual bandit / REINFORCE) that learns which prompt style, tool/retrieval policy, memory policy, planning depth, verification policy, and step budget to hand a frozen LLM for a given task -- without fine-tuning the model, and without letting the agent rewrite its own harness code (contrast Meta-Harness and HyperAgents).

Two papers come out of this repository (kept out of this public repo pending a possible double-blind workshop submission -- see below):

• Paper 1 -- the control-theoretic framing (context assembly as the controlled variable), with a stability and uncertainty-calibration analysis of the online controller, positioned against 2026 control-theoretic agent papers (Stable/Sparse Agentic Control, Regulatory Control Theory).
• Paper 2 -- the applied paper: the control system evaluated across three verifiable task domains and two model providers, the released dataset, and the deployment recipe.

See RECIPE.md for the org-facing deployment playbook independent of either paper.

Repository map

coe/                    core library
  types.py                Task/ContextConfig/AgentTrace/RewardBreakdown dataclasses
  context_space.py        the 729-configuration harness action space
  llm_adapter.py           dspy.LM adapters for Ollama / OpenAI / Bedrock
  cache.py                  SQLite response cache (not currently in the hot path -- see note below)
  dspy_programs.py         ContextConfig -> concrete DSPy Signature/Module assembly
  agent.py                 DSPyFrozenAgent: runs one (task, context) episode end-to-end
  domains/                 tool_use.py (synthetic CRM workflows), coding.py (HumanEval
                            subset), retrieval_qa.py (HotpotQA subset + local BM25)
  policies.py              random / fixed / epsilon-greedy bandit / REINFORCE
  reward.py                multi-objective reward decomposition
  memory.py                SQLite trace store + episodic-memory retrieval
scripts/
  build_datasets.py         materializes data/datasets/*.jsonl
  optimize_dspy_baseline.py  DSPy BootstrapFewShot static baseline per (domain, model)
  run_experiment.py          the optimizer x domain x model x seed matrix runner
  run_stability_experiment.py  supplementary run for Paper 1 (confidence logging)
  analyze_results.py         Paper 2 tables/figures
  analyze_stability.py       Paper 1 stability/calibration figures
configs/experiment.yaml     model roster, episode/seed budgets, cost cap
data/{datasets,runs,results}/
RECIPE.md                   the deployment playbook

Quick start

uv venv --python 3.12 .venv && source .venv/bin/activate
uv pip install -r requirements.txt   # or: dspy datasets boto3 numpy matplotlib pytest pandas scipy pyyaml python-dotenv

python scripts/build_datasets.py
python scripts/optimize_dspy_baseline.py
python scripts/run_experiment.py --config configs/experiment.yaml
python scripts/analyze_results.py

python scripts/run_stability_experiment.py
python scripts/analyze_stability.py

pytest

Provider credentials are read from a .env file one directory above this repo (see coe/llm_adapter.py::_load_env): OPENAI_API_KEY, AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / AWS_DEFAULT_REGION / BEDROCK_MODEL_HAIKU / BEDROCK_MODEL_SONNET. Ollama needs no key -- just a running local ollama serve with the model pulled (ollama pull qwen2.5:7b).

Dataset licensing

The coding domain is a subset of HumanEval (MIT license). The retrieval domain is a subset of HotpotQA (CC BY-SA 4.0) -- redistributed here as a small, attributed subset with its bundled gold+distractor paragraphs. The tool-use domain is fully synthetic (this repository's own CRM-workflow generator). The trajectory/reward logs generated by running the control system are original to this work.

A note on coe/cache.py

The SQLite response cache exists and is tested, but the hot path currently relies on dspy.LM(cache=True)'s own disk cache for de-duplicating identical (model, prompt) calls across episodes/optimizers -- which is what made re-running the DSPy static-baseline optimizer nearly instant on a second pass (Section 3.4 / Section 5 of Paper 2). coe/cache.py is there as the documented extension point for cost/usage accounting beyond what dspy.LM.history already tracks per call.