diff --git a/.gitignore b/.gitignore index 75e2da3b..80d3cefd 100644 --- a/.gitignore +++ b/.gitignore @@ -130,3 +130,6 @@ learning.db minion_logs/ *.oj-debug.json oj-debug.*.json + +# Generated orchestrator SFT/eval data artifacts +data/ diff --git a/scripts/orchestrator/build_orchestrator_sft.py b/scripts/orchestrator/build_orchestrator_sft.py new file mode 100644 index 00000000..b83c08a2 --- /dev/null +++ b/scripts/orchestrator/build_orchestrator_sft.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python +"""Build orchestrator SFT cold-start data by **base Qwen3-8B self-sampling** (v1). + +The orchestrator *is* the local Qwen3-8B. We roll the base (un-trained) model out +over the reasoning-SFT task pool (``load_sft_tasks``: GeneralThought + OpenThoughts +code/math/science), verify each trajectory's final answer against the gold +(``verify.make_verifier`` — math/code checkers + Gemini judge fallback), and keep +the cheapest-correct trajectory per task. Those passing trajectories become the +``conversations`` JSONL the SFT trainer consumes — i.e. the model learns from its +own successful rollouts (rejection-sampling cold-start, STaR-style). + +This is **v1** (base self-sampling): point ``--orchestrator-endpoint`` at the vLLM +server hosting the *base* ``Qwen/Qwen3-8B``. + +For **v2**, train on the v1 data, serve the v1 checkpoint with vLLM, then re-run +this exact driver with ``--orchestrator-endpoint`` pointed at the v1 checkpoint's +endpoint, ``--orchestrator-model`` set to the checkpoint name, and +``--samples-per-task 8`` — the (now stronger) policy self-samples a v2 dataset. + +The math/coder specialist tools are only wired when ``--math-endpoint`` / +``--coder-endpoint`` are passed (those are separately-served vLLM specialists); +otherwise they're omitted so the orchestrator only sees the tools it can actually +call. + +Example (v1, base self-sampling): + .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \ + --out data/orchestrator_sft_v1.jsonl \ + --orchestrator-endpoint http://localhost:8001/v1 \ + --orchestrator-model qwen3-8b \ + --samples-per-task 8 --max-keep-per-task 1 + +Example (v2, re-point at the v1 checkpoint): + .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \ + --out data/orchestrator_sft_v2.jsonl \ + --orchestrator-endpoint http://localhost:8010/v1 \ + --orchestrator-model orchestrator-sft-v1 \ + --samples-per-task 8 +""" + +from __future__ import annotations + +import argparse +import json +import logging +from typing import Optional + +from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog +from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout +from openjarvis.agents.hybrid.toolorchestra.unified import ( + make_call_orchestrator, + make_dispatch, +) +from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import ( + load_sft_tasks, +) +from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import ( + generate_sft_dataset, +) +from openjarvis.learning.intelligence.orchestrator.sft_data.verify import make_verifier + + +def main(argv: Optional[list[str]] = None) -> int: + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--out", default="data/orchestrator_sft_v1.jsonl") + # The orchestrator == the local Qwen3-8B self-sampling over its own rollouts. + p.add_argument("--orchestrator-endpoint", default="http://localhost:8001/v1", + help="OpenAI-compatible vLLM base URL serving the orchestrator.") + p.add_argument("--orchestrator-model", default="qwen3-8b") + p.add_argument("--orchestrator-api-key", default="EMPTY") + # Local OSS model endpoints; repeatable "model_id=base_url" (e.g. + # "Qwen/Qwen3.5-9B=http://localhost:8001/v1"). Unmapped local models are + # still listed but served unconfigured (base_url=None). + p.add_argument("--local-endpoint", action="append", default=[], + metavar="MODEL_ID=URL", + help="Local model id -> vLLM base URL (repeatable).") + p.add_argument("--max-tasks", type=int, default=None, + help="Cap on tasks (default: all of load_sft_tasks()).") + p.add_argument("--samples-per-task", type=int, default=8) + p.add_argument("--max-keep-per-task", type=int, default=1) + p.add_argument("--max-turns", type=int, default=8) + p.add_argument("--temperature", type=float, default=0.7) + args = p.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(message)s") + + local_endpoints = {} + for item in args.local_endpoint: + model_id, _, url = item.partition("=") + if model_id and url: + local_endpoints[model_id] = url + tools = orchestrator_catalog(local_endpoints=local_endpoints or None) + logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools]) + + call_orch = make_call_orchestrator( + args.orchestrator_model, + base_url=args.orchestrator_endpoint, + api_key=args.orchestrator_api_key, + temperature=args.temperature, + ) + dispatch = make_dispatch({}) + + def rollout_fn(task): + try: + return run_unified_rollout( + task.instruction, tools, + call_orchestrator=call_orch, dispatch=dispatch, + max_turns=args.max_turns, + ) + except Exception as exc: # network/key failures shouldn't kill the run + logging.warning("rollout failed for %s: %s", task.task_id, exc) + return None + + tasks = load_sft_tasks() + if args.max_tasks is not None: + tasks = tasks[: args.max_tasks] + logging.info("Loaded %d SFT tasks", len(tasks)) + + stats = generate_sft_dataset( + args.out, + tasks=tasks, + tools=tools, + rollout_fn=rollout_fn, + verify_fn=make_verifier(), + samples_per_task=args.samples_per_task, + max_keep_per_task=args.max_keep_per_task, + reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets highest reward + ) + print(json.dumps(stats, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/orchestrator/build_sft_data.py b/scripts/orchestrator/build_sft_data.py deleted file mode 100644 index 41ef73c7..00000000 --- a/scripts/orchestrator/build_sft_data.py +++ /dev/null @@ -1,19 +0,0 @@ -#!/usr/bin/env python -"""Thin CLI wrapper: build the orchestrator SFT dataset from NeuLab ADP. - -Equivalent to:: - - python -m openjarvis.learning.intelligence.orchestrator.sft_data.build [...] - -Usage: - python scripts/orchestrator/build_sft_data.py \ - --out data/orchestrator_sft_traces.jsonl \ - --max-tasks 2000 --adp-configs codeactinstruct,code_feedback,openhands -""" - -from __future__ import annotations - -from openjarvis.learning.intelligence.orchestrator.sft_data.build import _main - -if __name__ == "__main__": - raise SystemExit(_main()) diff --git a/scripts/orchestrator/build_unified_sft.py b/scripts/orchestrator/build_unified_sft.py index e4eea00b..40dff290 100644 --- a/scripts/orchestrator/build_unified_sft.py +++ b/scripts/orchestrator/build_unified_sft.py @@ -40,9 +40,9 @@ from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import ( def main(argv: Optional[list[str]] = None) -> int: p = argparse.ArgumentParser(description=__doc__) p.add_argument("--out", default="data/orchestrator_unified_sft.jsonl") - p.add_argument("--teacher-model", default="gpt-5") - p.add_argument("--teacher-base-url", default=None, - help="OpenAI-compatible base URL; omit for OpenAI cloud.") + p.add_argument("--teacher-model", default="claude-sonnet-4-6") + p.add_argument("--teacher-base-url", default="https://api.anthropic.com/v1/", + help="OpenAI-compatible base URL; pass '' for OpenAI cloud.") p.add_argument("--max-tasks", type=int, default=200) p.add_argument("--samples-per-task", type=int, default=4) p.add_argument("--max-keep-per-task", type=int, default=1) @@ -59,10 +59,16 @@ def main(argv: Optional[list[str]] = None) -> int: ) logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools]) + base_url = args.teacher_base_url or None + api_key = ( + os.environ.get("ANTHROPIC_API_KEY") + if base_url and "anthropic" in base_url + else os.environ.get("OPENAI_API_KEY") + ) call_orch = make_call_orchestrator( args.teacher_model, - base_url=args.teacher_base_url, - api_key=os.environ.get("OPENAI_API_KEY"), + base_url=base_url, + api_key=api_key, temperature=args.temperature, ) dispatch = make_dispatch({}) diff --git a/scripts/orchestrator/eval_orchestrator.py b/scripts/orchestrator/eval_orchestrator.py new file mode 100644 index 00000000..3e793387 --- /dev/null +++ b/scripts/orchestrator/eval_orchestrator.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python +"""Run the ToolOrchestra orchestrator over N-sample subsets of several +benchmarks and score them with the existing eval infra. + +Reuses, verbatim: + * ``openjarvis.evals.cli._build_dataset`` / ``_build_scorer`` / + ``_build_judge_backend`` (loaders + scorers + judge wiring), + * ``openjarvis.evals.core.runner.EvalRunner`` (parallel run + scoring), + * ``openjarvis.evals.core.types.RunConfig`` (run config), + +and plugs in our orchestrator as the "model" via +``openjarvis.learning.intelligence.orchestrator.eval_backend.OrchestratorBackend``. + +For each benchmark it runs ``EvalRunner(config, dataset, backend, scorer).run()``, +collects accuracy / cost / latency, writes a combined ``summary.json``, and +prints a table. + +NOTE on the judge: the OpenAI key is dead, so the default judge model is a +Gemini model (``gemini-2.5-flash``). The judge backend is still built via the +``cloud`` engine (``_build_judge_backend``) — the model string selects the +provider/route. Override with ``--judge-model`` / ``--judge-engine``. + +Example +------- + .venv/bin/python scripts/orchestrator/eval_orchestrator.py \\ + --benchmarks gaia,mmlu_pro --n 100 \\ + --orchestrator-endpoint http://localhost:8001/v1 \\ + --orchestrator-model qwen3-8b \\ + --output-dir ~/.openjarvis/experiments/hybrid/runs/orch-eval +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +# Friendly aliases (underscored / shorthand) -> the EXACT registry keys used by +# openjarvis.evals.cli.BENCHMARKS / _build_dataset / _build_scorer. +BENCHMARK_ALIASES: Dict[str, str] = { + "terminalbench_v2_1": "terminalbench-v2.1", + "terminalbench-v2_1": "terminalbench-v2.1", + "terminalbench_v21": "terminalbench-v2.1", + "terminalbench-v2.1": "terminalbench-v2.1", + "mmlu_pro": "mmlu-pro", + "mmlu-pro": "mmlu-pro", + "gaia": "gaia", + "taubench": "taubench", + "supergpqa": "supergpqa", +} + +DEFAULT_BENCHMARKS = "gaia,terminalbench_v2_1,taubench,mmlu_pro,supergpqa" + + +def _normalize_benchmark(name: str) -> str: + key = name.strip() + return BENCHMARK_ALIASES.get(key, BENCHMARK_ALIASES.get(key.lower(), key)) + + +def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Run the orchestrator over benchmark subsets and score them.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "--benchmarks", + default=DEFAULT_BENCHMARKS, + help="Comma-separated benchmark names (aliases normalized to registry keys).", + ) + p.add_argument("--n", type=int, default=100, help="Samples per benchmark.") + p.add_argument("--seed", type=int, default=42, help="Subset seed.") + p.add_argument( + "--orchestrator-endpoint", + default="http://localhost:8001/v1", + help="OpenAI-compatible base URL for the served orchestrator.", + ) + p.add_argument( + "--orchestrator-model", + default="qwen3-8b", + help="Served orchestrator model id.", + ) + p.add_argument( + "--orchestrator-api-key", + default="EMPTY", + help="API key for the orchestrator endpoint (EMPTY for local vLLM).", + ) + p.add_argument( + "--local-endpoint", + action="append", + default=[], + metavar="MODEL_ID=URL", + help=( + "Wire a local OSS tool: map a catalog model id to its vLLM base " + "URL, e.g. 'Qwen/Qwen3.5-9B=http://localhost:8003/v1'. Repeatable." + ), + ) + p.add_argument( + "--max-turns", + type=int, + default=8, + help="Max orchestrator turns per sample.", + ) + p.add_argument( + "--temperature", + type=float, + default=1.0, + help="Orchestrator sampling temperature.", + ) + p.add_argument( + "--max-workers", + type=int, + default=4, + help="Parallel samples per benchmark.", + ) + p.add_argument( + "--output-dir", + default="results/orchestrator-eval", + help="Directory for per-benchmark JSONL + combined summary.json.", + ) + # NOTE: OpenAI key is dead — default the judge to a Gemini model. + p.add_argument( + "--judge-model", + default="gemini-2.5-flash", + help="LLM-judge model (default Gemini — OpenAI key is dead).", + ) + p.add_argument( + "--judge-engine", + default="cloud", + help="Engine key for the judge backend.", + ) + return p.parse_args(argv) + + +def _run_one( + benchmark: str, + *, + n: int, + seed: int, + backend, + judge_model: str, + judge_engine: str, + max_workers: int, + orchestrator_model: str, + output_dir: Path, +) -> Dict[str, Any]: + """Run + score one benchmark, returning a row dict (or an error row).""" + # Imports here so --help works without importing the (heavy) eval stack. + from openjarvis.evals.cli import ( + _build_dataset, + _build_judge_backend, + _build_scorer, + ) + from openjarvis.evals.core.runner import EvalRunner + from openjarvis.evals.core.types import RunConfig + + output_path = output_dir / f"{benchmark.replace('.', '_')}_orchestrator.jsonl" + + dataset = _build_dataset(benchmark) + # The runner reloads internally; we also load here per spec so callers can + # inspect/size the subset up front. load() is idempotent for these sets. + dataset.load(max_samples=n, seed=seed) + + judge_backend = _build_judge_backend(judge_model, engine_key=judge_engine) + scorer = _build_scorer(benchmark, judge_backend, judge_model) + + config = RunConfig( + benchmark=benchmark, + backend="orchestrator", + model=orchestrator_model, + max_samples=n, + max_workers=max_workers, + seed=seed, + judge_model=judge_model, + judge_engine=judge_engine, + output_path=str(output_path), + ) + + runner = EvalRunner(config, dataset, backend, scorer) + started = time.time() + try: + summary = runner.run() + finally: + if judge_backend is not None: + judge_backend.close() + elapsed = time.time() - started + + return { + "benchmark": benchmark, + "accuracy": summary.accuracy, + "scored_samples": summary.scored_samples, + "correct": summary.correct, + "errors": summary.errors, + "total_samples": summary.total_samples, + "mean_latency_seconds": summary.mean_latency_seconds, + "total_cost_usd": summary.total_cost_usd, + "mean_continuous_score": summary.mean_continuous_score, + "wall_seconds": round(elapsed, 2), + "output_path": str(output_path), + } + + +def _print_table(rows: List[Dict[str, Any]]) -> None: + cols = [ + ("benchmark", 22, "{}"), + ("accuracy", 9, "{:.4f}"), + ("correct", 8, "{}"), + ("scored", 7, "{}"), + ("errors", 7, "{}"), + ("cost($)", 10, "{:.4f}"), + ("lat(s)", 9, "{:.2f}"), + ("wall(s)", 9, "{:.1f}"), + ] + header = " ".join(f"{name:<{w}}" for name, w, _ in cols) + print("\n" + header) + print("-" * len(header)) + for r in rows: + if r.get("error"): + print(f"{r['benchmark']:<22} ERROR: {r['error']}") + continue + vals = { + "benchmark": r["benchmark"], + "accuracy": r["accuracy"], + "correct": r["correct"], + "scored": r["scored_samples"], + "errors": r["errors"], + "cost($)": r["total_cost_usd"], + "lat(s)": r["mean_latency_seconds"], + "wall(s)": r["wall_seconds"], + } + cells = [] + for name, w, fmt in cols: + v = vals[name] + try: + s = fmt.format(v) + except (ValueError, TypeError): + s = str(v) + cells.append(f"{s:<{w}}") + print(" ".join(cells)) + + +def main(argv: Optional[List[str]] = None) -> int: + args = _parse_args(argv) + + benchmarks = [ + _normalize_benchmark(b) for b in args.benchmarks.split(",") if b.strip() + ] + output_dir = Path(args.output_dir).expanduser() + output_dir.mkdir(parents=True, exist_ok=True) + + # Build the orchestrator backend once (stateless across benchmarks). + from openjarvis.learning.intelligence.orchestrator.eval_backend import ( + OrchestratorBackend, + ) + + local_endpoints: Dict[str, str] = {} + for pair in args.local_endpoint: + if "=" not in pair: + raise SystemExit( + f"--local-endpoint expects MODEL_ID=URL, got: {pair!r}" + ) + model_id, url = pair.split("=", 1) + local_endpoints[model_id.strip()] = url.strip() + + backend = OrchestratorBackend( + orchestrator_endpoint=args.orchestrator_endpoint, + orchestrator_model=args.orchestrator_model, + api_key=args.orchestrator_api_key, + local_endpoints=local_endpoints, + max_turns=args.max_turns, + temperature=args.temperature, + ) + + rows: List[Dict[str, Any]] = [] + try: + for benchmark in benchmarks: + print(f"\n=== {benchmark} (n={args.n}, seed={args.seed}) ===") + try: + row = _run_one( + benchmark, + n=args.n, + seed=args.seed, + backend=backend, + judge_model=args.judge_model, + judge_engine=args.judge_engine, + max_workers=args.max_workers, + orchestrator_model=args.orchestrator_model, + output_dir=output_dir, + ) + print( + f" accuracy={row['accuracy']:.4f} " + f"({row['correct']}/{row['scored_samples']}) " + f"errors={row['errors']} cost=${row['total_cost_usd']:.4f}" + ) + except Exception as exc: # noqa: BLE001 - one bad bench shouldn't kill all + print(f" FAILED: {type(exc).__name__}: {exc}", file=sys.stderr) + row = {"benchmark": benchmark, "error": f"{type(exc).__name__}: {exc}"} + rows.append(row) + finally: + backend.close() + + combined = { + "orchestrator_model": args.orchestrator_model, + "orchestrator_endpoint": args.orchestrator_endpoint, + "judge_model": args.judge_model, + "judge_engine": args.judge_engine, + "n": args.n, + "seed": args.seed, + "max_turns": args.max_turns, + "benchmarks": rows, + } + summary_path = output_dir / "summary.json" + with open(summary_path, "w") as f: + json.dump(combined, f, indent=2, default=str) + + _print_table(rows) + print(f"\nCombined summary written to {summary_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/openjarvis/agents/hybrid/_prices.py b/src/openjarvis/agents/hybrid/_prices.py index 1c8c546b..c2dcae45 100644 --- a/src/openjarvis/agents/hybrid/_prices.py +++ b/src/openjarvis/agents/hybrid/_prices.py @@ -11,6 +11,7 @@ from __future__ import annotations # USD per million tokens, (input, output). Local models = 0. PRICES: dict[str, tuple[float, float]] = { + "claude-opus-4-8": (5.00, 25.0), "claude-opus-4-7": (5.00, 25.0), "claude-sonnet-4-6": (3.00, 15.0), "claude-haiku-4-5": (1.00, 5.00), diff --git a/src/openjarvis/agents/hybrid/expert_registry.py b/src/openjarvis/agents/hybrid/expert_registry.py index 703b2193..344ebd6d 100644 --- a/src/openjarvis/agents/hybrid/expert_registry.py +++ b/src/openjarvis/agents/hybrid/expert_registry.py @@ -24,6 +24,7 @@ circular import. from __future__ import annotations import random +import re from dataclasses import dataclass, field from typing import Dict, List, Optional @@ -34,13 +35,30 @@ KIND_MODEL = "model" # an LLM exposed as a tool (the paper's "models as tools") KIND_WEB_SEARCH = "web_search" KIND_LOCAL_SEARCH = "local_search" KIND_CODE = "code_interpreter" +KIND_TOOL = "tool" # a bridged real OpenJarvis tool (custom param schema) -VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_LOCAL_SEARCH, KIND_CODE) +VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_LOCAL_SEARCH, KIND_CODE, KIND_TOOL) -# Backend dispatch types understood by ``toolorchestra._call_worker``. +# Flat-catalog category label, surfaced in each tool's spec so the orchestrator +# can tell tool types apart without us imposing any hierarchy (the menu stays +# flat — this is just a tag). +CATEGORY_GENERALIST = "generalist_model" +CATEGORY_SPECIALIZED = "specialized_model" +CATEGORY_BASIC = "basic_tool" +CATEGORY_SMALL_GENERALIST = "small_generalist" +CATEGORY_STRONG_GENERALIST = "strong_generalist" +# Two-model-class taxonomy for the orchestrator catalog (the specialized / +# small-vs-strong generalist tiers are superseded). +CATEGORY_CLOUD_FRONTIER = "cloud_frontier" +CATEGORY_LOCAL_OSS = "local_open_source" + +# Backend dispatch types understood by ``toolorchestra._call_worker`` (plus the +# ``openjarvis-tool`` bridge, dispatched in ``unified.make_dispatch`` via the +# OpenJarvis ToolExecutor rather than ``_call_worker``). VALID_BACKENDS = ( "vllm", "openai", "anthropic", "gemini", "openrouter", "anthropic-web-search", "tavily-search", "modal-python", + "openjarvis-tool", ) @@ -63,6 +81,11 @@ class ExpertTool: price_in: float = 0.0 price_out: float = 0.0 latency_s: float = 5.0 + category: str = "" # generalist_model | specialized_model | basic_tool + # Optional explicit JSON-schema for the tool's arguments. Set for bridged + # real OpenJarvis tools (``openjarvis-tool`` backend) whose params don't fit + # the fixed kind-based schemas; takes precedence over the kind default. + param_schema: Optional[dict] = None def __post_init__(self) -> None: if self.kind not in VALID_KINDS: @@ -75,7 +98,13 @@ class ExpertTool: # ---- orchestrator-visible spec ------------------------------------- def _param_schema(self) -> Dict[str, object]: - """JSON-schema for the tool's arguments (one typed param per kind).""" + """JSON-schema for the tool's arguments (one typed param per kind). + + An explicit ``param_schema`` (set by :func:`openjarvis_tool` for bridged + real tools) overrides the kind-based default. + """ + if self.param_schema is not None: + return self.param_schema if self.kind == KIND_WEB_SEARCH or self.kind == KIND_LOCAL_SEARCH: return { "type": "object", @@ -122,21 +151,36 @@ class ExpertTool: return self.summary.rstrip(".") + "." + cost_line def to_spec(self) -> Dict[str, object]: - """OpenAI-style tool spec the orchestrator conditions on.""" - return { - "type": "function", - "function": { - "name": self.name, - "description": self.description(), - "parameters": self._param_schema(), - }, + """OpenAI-style tool spec the orchestrator conditions on. + + Flat list, but each function carries a ``category`` tag so the policy can + distinguish generalist vs specialized models vs basic tools. + """ + fn: Dict[str, object] = { + "name": self.name, + "description": self.description(), + "parameters": self._param_schema(), } + if self.category: + fn["category"] = self.category + return {"type": "function", "function": fn} def _price(model: str) -> tuple[float, float]: return PRICES.get(model, (0.0, 0.0)) +def _tool_name(model: str) -> str: + """Tool-safe function name derived from a model id (``qwen3-8b`` -> ``qwen3_8b``). + + Strips any provider prefix and replaces non-alphanumerics with underscores so + the catalog exposes one named tool per concrete model. + """ + base = model.split("/")[-1].lower() + safe = re.sub(r"[^a-z0-9]+", "_", base).strip("_") + return safe or "local_model" + + # Default catalog: the paper's tool categories, mapped onto the models/tools # OpenJarvis can actually call. One named tool per model (faithful §3.1). def default_catalog( @@ -151,70 +195,300 @@ def default_catalog( """ cat: List[ExpertTool] = [] - # ---- generalist / frontier models ---- - for name, model, summary, lat in [ - ("gpt_5", "gpt-5", + # ---- generalist / frontier models (one named tool per model VERSION) ---- + for model, summary, lat in [ + ("gpt-5", "Frontier generalist (GPT-5). Strongest reasoning across domains.", 30.0), - ("gpt_5_mini", "gpt-5-mini", + ("gpt-5-mini", "Mid-tier generalist (GPT-5-mini). Solid reasoning, much cheaper.", 15.0), - ("gpt_4o", "gpt-4o", + ("gpt-4o", "Fast generalist (GPT-4o). Good for simple steps and formatting.", 8.0), - ("claude_opus", "claude-opus-4-7", - "Frontier generalist (Claude Opus). Strong long-horizon reasoning.", 26.0), - ("claude_sonnet", "claude-sonnet-4-6", - "Strong generalist (Claude Sonnet). Balanced cost/capability.", 15.0), - ("gemini_2_5_pro", "gemini-2.5-pro", + ("claude-opus-4-7", + "Frontier generalist (Claude Opus 4.7). Strong long-horizon reasoning.", 26.0), + ("claude-sonnet-4-6", + "Strong generalist (Claude Sonnet 4.6). Balanced cost/capability.", 15.0), + ("gemini-2.5-pro", "Frontier generalist (Gemini 2.5 Pro). Strong multimodal reasoning.", 20.0), - ("gemini_2_5_flash", "gemini-2.5-flash", + ("gemini-2.5-flash", "Cheap fast generalist (Gemini 2.5 Flash).", 8.0), - ("llama_3_3_70b", "meta-llama/llama-3.3-70b-instruct", + ("meta-llama/llama-3.3-70b-instruct", "Open generalist (Llama-3.3-70B). Decent general knowledge, low cost.", 10.0), - ("qwen3_32b", "qwen/qwen3-32b", + ("qwen/qwen3-32b", "Open generalist (Qwen3-32B). Strong math/science reasoning, low cost.", 9.0), ]: - ep = "openai" if name.startswith("gpt") else ( - "anthropic" if name.startswith("claude") else ( - "gemini" if name.startswith("gemini") else "openrouter")) + ep = ("openai" if model.startswith("gpt") else + "anthropic" if model.startswith("claude") else + "gemini" if model.startswith("gemini") else "openrouter") pi, po = _price(model) cat.append(ExpertTool( - name=name, kind=KIND_MODEL, backend_type=ep, summary=summary, + name=_tool_name(model), kind=KIND_MODEL, backend_type=ep, summary=summary, model=model, price_in=pi, price_out=po, latency_s=lat, + category=CATEGORY_GENERALIST, )) # ---- specialized: code ---- - pi, po = _price("qwen/qwen-2.5-coder-32b-instruct") + coder = "qwen/qwen-2.5-coder-32b-instruct" + pi, po = _price(coder) cat.append(ExpertTool( - name="qwen2_5_coder_32b", kind=KIND_MODEL, backend_type="openrouter", + name=_tool_name(coder), kind=KIND_MODEL, backend_type="openrouter", summary="Specialized code model (Qwen2.5-Coder-32B). Writes/debugs code.", - model="qwen/qwen-2.5-coder-32b-instruct", - price_in=pi, price_out=po, latency_s=9.0, + model=coder, price_in=pi, price_out=po, latency_s=9.0, + category=CATEGORY_SPECIALIZED, )) # ---- local backbone as a tool (on-device vLLM), if served ---- + # Named after the actual served model (faithful "one named tool per model"), + # not a generic "local_model" — e.g. "qwen3-8b" -> tool "qwen3_8b". if local_model and local_endpoint: cat.append(ExpertTool( - name="local_model", kind=KIND_MODEL, backend_type="vllm", - summary=("On-device open model served locally. Cheap and private; " - "good for extraction, formatting, arithmetic on given data."), + name=_tool_name(local_model), kind=KIND_MODEL, backend_type="vllm", + summary=(f"On-device open model ({local_model}) served locally. Cheap " + "and private; good for extraction, formatting, arithmetic on " + "given data."), model=local_model, base_url=local_endpoint, price_in=0.0, price_out=0.0, latency_s=2.0, + category=CATEGORY_GENERALIST, )) # ---- basic tools ---- cat.append(ExpertTool( name="web_search", kind=KIND_WEB_SEARCH, backend_type="tavily-search", summary="Web search (Tavily). Use for facts that need a live lookup.", - model="tavily", latency_s=8.0, + model="tavily", latency_s=8.0, category=CATEGORY_BASIC, )) cat.append(ExpertTool( name="code_interpreter", kind=KIND_CODE, backend_type="modal-python", summary="Python sandbox. Execute code and return stdout/stderr.", - model="modal-python", latency_s=6.0, + model="modal-python", latency_s=6.0, category=CATEGORY_BASIC, )) return cat +def openjarvis_tool( + registered_name: str, + *, + summary: str, + params: dict, + latency_s: float = 5.0, +) -> ExpertTool: + """Build an :class:`ExpertTool` that bridges a real OpenJarvis tool. + + ``registered_name`` is the tool's key in ``ToolRegistry`` (e.g. ``calculator``, + ``shell_exec``). ``params`` is the JSON-schema *properties*-style dict for the + tool's arguments; it is surfaced verbatim by :meth:`ExpertTool.to_spec`. The + resulting tool dispatches through the OpenJarvis ``ToolExecutor`` (backend + ``openjarvis-tool``) rather than ``_call_worker``. + """ + return ExpertTool( + name=registered_name, + kind=KIND_TOOL, + backend_type="openjarvis-tool", + summary=summary, + model=registered_name, + latency_s=latency_s, + category=CATEGORY_BASIC, + param_schema=params, + ) + + +# Real OpenJarvis tools bridged into the orchestrator catalog as basic tools. +# Names must match the ``ToolRegistry`` keys (confirmed present: calculator, +# shell_exec, file_read, file_write, http_request). +def _openjarvis_basic_tools() -> List[ExpertTool]: + def obj(properties: dict, required: List[str]) -> dict: + return {"type": "object", "properties": properties, "required": required} + + return [ + openjarvis_tool( + "calculator", + summary="Evaluate an arithmetic / math expression and return the result.", + params=obj( + {"expression": {"type": "string", + "description": "Math expression to evaluate."}}, + ["expression"], + ), + latency_s=1.0, + ), + openjarvis_tool( + "shell_exec", + summary=("Run a shell command and return its stdout/stderr. Critical " + "for terminal / TerminalBench-style tasks."), + params=obj( + {"command": {"type": "string", + "description": "Shell command to execute."}}, + ["command"], + ), + latency_s=4.0, + ), + openjarvis_tool( + "file_read", + summary="Read the contents of a file at the given path.", + params=obj( + {"path": {"type": "string", "description": "Path of the file to read."}}, + ["path"], + ), + latency_s=1.0, + ), + openjarvis_tool( + "file_write", + summary="Write content to a file at the given path.", + params=obj( + {"path": {"type": "string", "description": "Path of the file to write."}, + "content": {"type": "string", "description": "Content to write."}}, + ["path", "content"], + ), + latency_s=1.0, + ), + openjarvis_tool( + "http_request", + summary="Make an HTTP request to a URL and return the response body.", + params={ + "type": "object", + "properties": { + "url": {"type": "string", "description": "Request URL."}, + "method": {"type": "string", + "description": "HTTP method (GET, POST, ...). Default GET."}, + }, + "required": ["url"], + }, + latency_s=4.0, + ), + openjarvis_tool( + "think", + summary=("Record a private reasoning step (scratchpad). No external " + "effect; use to plan before acting on hard reasoning tasks."), + params=obj( + {"thought": {"type": "string", + "description": "Your reasoning or thought process."}}, + ["thought"], + ), + latency_s=0.5, + ), + openjarvis_tool( + "apply_patch", + summary=("Apply a unified-diff patch to a file. Use to edit code for " + "terminal / SWE-style tasks."), + params=obj( + {"patch": {"type": "string", + "description": "The unified diff patch text to apply."}, + "path": {"type": "string", + "description": "Target file path (auto-detected from the " + "patch header if omitted)."}}, + ["patch"], + ), + latency_s=2.0, + ), + openjarvis_tool( + "pdf_extract", + summary=("Extract text from a PDF file. Use for GAIA-style tasks with " + "PDF attachments."), + params=obj( + {"file_path": {"type": "string", + "description": "Path to the PDF file."}, + "pages": {"type": "string", + "description": "Page range, e.g. '1-5' or '1,3,5'. " + "Omit for all pages."}}, + ["file_path"], + ), + latency_s=3.0, + ), + openjarvis_tool( + "db_query", + summary=("Run a SQL query against a SQLite/Postgres database and return " + "rows. Read-only by default."), + params=obj( + {"query": {"type": "string", + "description": "SQL query to execute."}, + "db_path": {"type": "string", + "description": "Path to a SQLite DB file. Defaults to " + "in-memory."}, + "read_only": {"type": "boolean", + "description": "Restrict to SELECT/EXPLAIN/PRAGMA. " + "Default: true."}}, + ["query"], + ), + latency_s=3.0, + ), + ] + + +# Local-Cloud Hybrid orchestrator catalog. The menu is two model classes — cloud +# frontier models and locally-served open-source models — plus the basic tools +# (web search, code interpreter, and the bridged real OpenJarvis tools). The +# orchestrator model itself is NOT in the catalog. +# +# Local model ids -> vLLM base_url come from ``local_endpoints``; unmapped ids +# get ``base_url=None`` (tool still listed, dispatch wires the endpoint later). +_LOCAL_OSS_MODELS = ( + "Qwen/Qwen3.5-9B", + "Qwen/Qwen3.6-27B-FP8", + "Qwen/Qwen3.5-122B-A10B-FP8", + "Qwen/Qwen3.5-397B-A17B-FP8", +) + +_CLOUD_FRONTIER_MODELS = ( + # (model, backend, summary, latency_s) + ("gpt-5.5", "openai", + "Cloud frontier generalist (GPT-5.5). Strongest reasoning across domains.", 30.0), + ("claude-opus-4-8", "anthropic", + "Cloud frontier generalist (Claude Opus 4.8). Strong long-horizon reasoning.", 26.0), +) + + +def orchestrator_catalog( + *, + local_endpoints: Optional[Dict[str, str]] = None, + include_tools: bool = True, +) -> List[ExpertTool]: + """Return the orchestrator's tool catalog: two model classes + basic tools. + + ``local_endpoints`` maps a local model id (e.g. ``"Qwen/Qwen3.5-9B"``) to its + vLLM ``base_url``; ids not present get ``base_url=None``. ``include_tools`` + (default True) appends the basic tools — web search, code interpreter, and the + bridged real OpenJarvis tools (calculator, shell_exec, file_read, file_write, + http_request). + """ + local_endpoints = local_endpoints or {} + cat: List[ExpertTool] = [] + + # ---- cloud frontier models ---- + for model, backend, summary, lat in _CLOUD_FRONTIER_MODELS: + pi, po = _price(model) + cat.append(ExpertTool( + name=_tool_name(model), kind=KIND_MODEL, backend_type=backend, + summary=summary, model=model, price_in=pi, price_out=po, + latency_s=lat, category=CATEGORY_CLOUD_FRONTIER, + )) + + # ---- local open-source models (served via vLLM, price 0/0) ---- + for model in _LOCAL_OSS_MODELS: + cat.append(ExpertTool( + name=_tool_name(model), kind=KIND_MODEL, backend_type="vllm", + summary=(f"Locally-served open-source model ({model}). Cheap and " + "private; no per-token cost."), + model=model, base_url=local_endpoints.get(model), + price_in=0.0, price_out=0.0, latency_s=4.0, + category=CATEGORY_LOCAL_OSS, + )) + + if include_tools: + # ---- basic tools ---- + cat.append(ExpertTool( + name="web_search", kind=KIND_WEB_SEARCH, backend_type="tavily-search", + summary="Web search (Tavily). Use for facts that need a live lookup.", + model="tavily", latency_s=8.0, category=CATEGORY_BASIC, + )) + cat.append(ExpertTool( + name="code_interpreter", kind=KIND_CODE, backend_type="modal-python", + summary="Python sandbox. Execute code and return stdout/stderr.", + model="modal-python", latency_s=6.0, category=CATEGORY_BASIC, + )) + cat.extend(_openjarvis_basic_tools()) + + return cat + + def build_tool_specs(tools: List[ExpertTool]) -> List[Dict[str, object]]: """Turn a tool list into the OpenAI-style tools JSON the policy sees.""" return [t.to_spec() for t in tools] @@ -298,13 +572,23 @@ def to_worker_dict(tool: ExpertTool) -> Dict[str, object]: __all__ = [ + "CATEGORY_BASIC", + "CATEGORY_CLOUD_FRONTIER", + "CATEGORY_GENERALIST", + "CATEGORY_LOCAL_OSS", + "CATEGORY_SMALL_GENERALIST", + "CATEGORY_SPECIALIZED", + "CATEGORY_STRONG_GENERALIST", "ExpertTool", "KIND_CODE", "KIND_LOCAL_SEARCH", "KIND_MODEL", + "KIND_TOOL", "KIND_WEB_SEARCH", "build_tool_specs", "default_catalog", + "openjarvis_tool", + "orchestrator_catalog", "sample_tool_config", "to_worker_dict", "tools_by_name", diff --git a/src/openjarvis/agents/hybrid/toolorchestra/rollout.py b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py index a783f7dc..6140a0fb 100644 --- a/src/openjarvis/agents/hybrid/toolorchestra/rollout.py +++ b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py @@ -28,6 +28,27 @@ from openjarvis.agents.hybrid.expert_registry import ( RL_ORCHESTRATOR_SYS = "You are good at using tools." + +def build_system_prompt(specs: List[Dict[str, object]]) -> str: + """Faithful ToolOrchestra system prompt (arXiv:2511.21689, verbatim from + their ``prepare_sft_data.py``): the ``RL_ORCHESTRATOR_SYS`` preamble + the + Qwen-style ````/```` block. Deliberately contains NO + routing/delegation instructions — cost-aware routing is learned from the RL + reward, not prompted. + """ + tools_block = "\n".join(json.dumps(s) for s in specs) + return ( + f"{RL_ORCHESTRATOR_SYS}\n\n# Tools\n\n" + "You may call one or more functions to assist with the user query.\n\n" + "You are provided with function signatures within " + "XML tags:\n" + f"\n{tools_block}\n\n\n" + "For each function call, return a json object with function name and " + "arguments within XML tags:\n" + '\n{"name": , "arguments": }' + "\n" + ) + # Char-level cap on the accumulated context (mirrors the paper's ~24k-token cap). _CONTEXT_CAP = 24000 @@ -137,6 +158,7 @@ __all__ = [ "RL_ORCHESTRATOR_SYS", "UnifiedRollout", "UnifiedTurn", + "build_system_prompt", "run_unified_rollout", "tool_call_tag", ] diff --git a/src/openjarvis/agents/hybrid/toolorchestra/unified.py b/src/openjarvis/agents/hybrid/toolorchestra/unified.py index 4a861c0f..1626a9aa 100644 --- a/src/openjarvis/agents/hybrid/toolorchestra/unified.py +++ b/src/openjarvis/agents/hybrid/toolorchestra/unified.py @@ -10,6 +10,7 @@ network or keys. from __future__ import annotations +import json from typing import Any, Callable, Dict, List, Optional, Tuple from openjarvis.agents.hybrid._prices import cost as _model_cost @@ -63,13 +64,74 @@ def make_call_orchestrator( return call_orchestrator +def _dispatch_openjarvis_tool( + tool: ExpertTool, + arguments: Dict[str, Any], + executor_holder: Dict[str, Any], +) -> Tuple[str, float, int, bool]: + """Execute a bridged real OpenJarvis tool via its ``ToolExecutor``. + + Builds the executor lazily (cached in ``executor_holder``) and degrades to a + clear error string — never a crash — if the tool registry / executor can't be + instantiated. Returns ``(content, cost_usd, total_tokens, is_local=True)``. + """ + try: + executor = executor_holder.get("executor") + if executor is None: + from openjarvis.core.registry import ToolRegistry + from openjarvis.tools._stubs import ToolExecutor + + instances = [] + for name in ToolRegistry.keys(): + entry = ToolRegistry.get(name) + try: + instances.append(entry() if isinstance(entry, type) else entry) + except Exception: + continue + # Auto-approve confirmation-gated tools (shell_exec, git_commit, ...): + # rollouts/eval run headless in a sandbox, so there is no human to + # confirm. Without this, those tools return a "requires confirmation" + # error instead of executing — which would silently break TerminalBench. + executor = ToolExecutor( + instances, + interactive=True, + confirm_callback=lambda _prompt: True, + ) + executor_holder["executor"] = executor + + from openjarvis.core.types import ToolCall + + call = ToolCall( + id=f"orch-{tool.name}", + name=str(tool.model), + arguments=json.dumps(arguments or {}), + ) + result = executor.execute(call) + usage = getattr(result, "usage", None) or {} + total_tokens = int(usage.get("total_tokens", 0) or 0) + return ( + str(getattr(result, "content", "")), + float(getattr(result, "cost_usd", 0.0) or 0.0), + total_tokens, + True, + ) + except Exception as exc: # never crash the rollout on a tool-bridge failure + return (f"[openjarvis-tool error: {tool.model}: {exc}]", 0.0, 0, True) + + def make_dispatch( cfg: Optional[Dict[str, Any]] = None, ) -> Callable[[ExpertTool, Dict[str, Any]], Tuple[str, float, int, bool]]: """Tool-execution caller: run the chosen tool and return (obs, cost, tokens, is_local).""" cfg = cfg or {} + executor_holder: Dict[str, Any] = {} def dispatch(tool: ExpertTool, arguments: Dict[str, Any]): + # Bridged real OpenJarvis tools run through the ToolExecutor, not + # the model/worker path. + if tool.backend_type == "openjarvis-tool": + return _dispatch_openjarvis_tool(tool, arguments or {}, executor_holder) + worker = to_worker_dict(tool) prompt = "" for key in ("input", "query", "code"): diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-grpo.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-grpo.toml new file mode 100644 index 00000000..5f0c4be4 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-grpo.toml @@ -0,0 +1,54 @@ +# Orchestrator GRPO — Gemma-4-12B (cost-aware, lambda-swept). +# +# Trajectory-level GRPO over the UNIFIED tool-use rollout (the same faithful +# multi-turn path the SFT data is built from). For each prompt we +# sample group_size rollouts with the *policy* as orchestrator +# (run_unified_rollout), verify each rollout's final_answer against the gold +# answer, score with the cost-aware reward (r = +1/-1, R -= lam * cost / +# cost_max), normalise advantages within the group, and PPO-clip the +# trajectory log-prob update (summed over assistant-masked tokens only). +# +# Batch math: prompts_per_step * group_size = 256 * 8 = 2048 rollouts per +# step. grpo_max_prompts / prompts_per_step = 20000 / 256 ~= 78 batches +# per epoch. + +[model] +model_name = "google/gemma-4-12B-it" +max_prompt_length = 24000 +max_response_length = 8768 + +[training] +num_epochs = 5 +learning_rate = 1e-6 +max_grad_norm = 1.0 +gradient_checkpointing = true + +[grpo] +# group_size = number of trajectories sampled per prompt. +num_samples_per_prompt = 8 +temperature = 1.0 +kl_coef = 0.0001 +clip_ratio = 0.2 +# Per-rollout tool-use turn cap (one assistant turn per loop iteration). +max_turns = 10 + +[data] +grpo_max_prompts = 20000 # pooled GRPO prompts (cached on the trainer) +prompts_per_step = 256 # prompts per _grpo_step (256 * 8 = 2048 traj/step) + +[reward] +# Cost-aware reward: R = base(+1/-1) - lam * (cost_usd / cost_max). +# `lam` is SWEPT across runs to trace the accuracy/cost Pareto frontier, +# e.g. [0.0, 0.05, 0.1, 0.2, 0.4]. The value below is the default run. +lam = 0.0 +cost_max = 0.10 + +# Optional: map local OSS model ids -> served vLLM base_url so the local-model +# tools in the orchestrator catalog dispatch to a real endpoint. Omit to leave +# them unbacked (tool still listed). +[local_endpoints] +# "Qwen/Qwen3.5-9B" = "http://localhost:8001/v1" + +[checkpoint] +checkpoint_dir = "checkpoints/orchestrator_grpo_gemma_4_12b" +save_every_n_epochs = 1 diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-sft.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-sft.toml new file mode 100644 index 00000000..26fe2d3f --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-sft.toml @@ -0,0 +1,34 @@ +# Orchestrator: gemma-4-12B-it. v1 = base self-sampling (the served orchestrator_model is this same base). +# +# v1 data = base gemma-4-12B-it self-sampling: roll the base model out over +# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the +# dataset first (no GPU, needs the vLLM orchestrator endpoint up): +# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \ +# --out data/orchestrator_gemma-4-12b_sft_v1.jsonl --samples-per-task 8 +# then train on a rented GPU with regenerate_traces = false. + +[model] +model_name = "google/gemma-4-12B-it" +max_seq_length = 4096 + +[training] +num_epochs = 3 +batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch +learning_rate = 2e-5 +weight_decay = 0.01 +warmup_ratio = 0.1 +gradient_checkpointing = true + +[data] +trace_cache_path = "data/orchestrator_gemma-4-12b_sft_v1.jsonl" +regenerate_traces = false +# Base gemma-4-12B-it self-sampling over load_sft_tasks() (see sft_data/reject_sample.py). +orchestrator_endpoint = "http://localhost:8001/v1" +orchestrator_model = "gemma-4-12b" +samples_per_task = 8 +max_keep_per_task = 1 +distill_max_tasks = 0 # 0 -> all of load_sft_tasks() + +[checkpoint] +checkpoint_dir = "checkpoints/orchestrator_gemma-4-12b_sft" +save_every_n_epochs = 1 diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-grpo.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-grpo.toml new file mode 100644 index 00000000..35051898 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-grpo.toml @@ -0,0 +1,44 @@ +# Orchestrator GRPO — Qwen3-8B (cost-aware, lambda-swept). +# +# Fields map onto OrchestratorGRPOConfig. GRPO uses group-relative +# advantages (no critic): for each prompt, sample group_size trajectories, +# verify against the gold answer, score with the cost-aware reward +# (r = +1/-1, R -= lam * cost / cost_max), normalise advantages within the +# group, and PPO-clip the policy update. +# +# Batch math: prompts_per_step * group_size = 256 * 8 = 2048 trajectories +# per step. grpo_max_prompts / prompts_per_step = 20000 / 256 ~= 78 batches +# per epoch. + +[model] +model_name = "Qwen/Qwen3-8B" +max_prompt_length = 24000 +max_response_length = 8768 + +[training] +num_epochs = 5 +learning_rate = 1e-6 +max_grad_norm = 1.0 +gradient_checkpointing = true + +[grpo] +# group_size = number of trajectories sampled per prompt. +num_samples_per_prompt = 8 +temperature = 1.0 +kl_coef = 0.0001 +clip_ratio = 0.2 + +[data] +grpo_max_prompts = 20000 # pooled GRPO prompts (cached on the trainer) +prompts_per_step = 256 # prompts per _grpo_step (256 * 8 = 2048 traj/step) + +[reward] +# Cost-aware reward: R = base(+1/-1) - lam * (cost_usd / cost_max). +# `lam` is SWEPT across runs to trace the accuracy/cost Pareto frontier, +# e.g. [0.0, 0.05, 0.1, 0.2, 0.4]. The value below is the default run. +lam = 0.0 +cost_max = 0.10 + +[checkpoint] +checkpoint_dir = "checkpoints/orchestrator_grpo" +save_every_n_epochs = 1 diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml index 57d8581b..9553b587 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml @@ -1,8 +1,10 @@ # Orchestrator SFT cold-start — Qwen3-8B (~9B, "easier to train"). # -# Fields map onto OrchestratorSFTConfig. Build the dataset first (no GPU): -# python -m openjarvis.learning.intelligence.orchestrator.sft_data.build \ -# --out data/orchestrator_sft_traces.jsonl --max-tasks 2000 +# v1 data = base Qwen3-8B self-sampling: roll the base model out over +# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the +# dataset first (no GPU, needs the vLLM orchestrator endpoint up): +# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \ +# --out data/orchestrator_sft_v1.jsonl --samples-per-task 8 # then train on a rented GPU with regenerate_traces = false. [model] @@ -11,18 +13,21 @@ max_seq_length = 4096 [training] num_epochs = 3 -batch_size = 8 +batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch learning_rate = 2e-5 weight_decay = 0.01 warmup_ratio = 0.1 gradient_checkpointing = true [data] -trace_cache_path = "data/orchestrator_sft_traces.jsonl" +trace_cache_path = "data/orchestrator_sft_v1.jsonl" regenerate_traces = false -# Coding / agentic / tool-use ADP sub-configs (per the meeting steer). -adp_configs = "codeactinstruct,code_feedback,openhands,agenttuning_os,swe-smith" -distill_max_tasks = 2000 +# Base Qwen3-8B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py). +orchestrator_endpoint = "http://localhost:8001/v1" +orchestrator_model = "qwen3-8b" +samples_per_task = 8 +max_keep_per_task = 1 +distill_max_tasks = 0 # 0 -> all of load_sft_tasks() [checkpoint] checkpoint_dir = "checkpoints/orchestrator_sft" diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-grpo.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-grpo.toml new file mode 100644 index 00000000..7279bfd1 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-grpo.toml @@ -0,0 +1,54 @@ +# Orchestrator GRPO — Qwen3.5-9B (cost-aware, lambda-swept). +# +# Trajectory-level GRPO over the UNIFIED tool-use rollout (the same faithful +# multi-turn path the SFT data is built from). For each prompt we +# sample group_size rollouts with the *policy* as orchestrator +# (run_unified_rollout), verify each rollout's final_answer against the gold +# answer, score with the cost-aware reward (r = +1/-1, R -= lam * cost / +# cost_max), normalise advantages within the group, and PPO-clip the +# trajectory log-prob update (summed over assistant-masked tokens only). +# +# Batch math: prompts_per_step * group_size = 256 * 8 = 2048 rollouts per +# step. grpo_max_prompts / prompts_per_step = 20000 / 256 ~= 78 batches +# per epoch. + +[model] +model_name = "Qwen/Qwen3.5-9B" +max_prompt_length = 24000 +max_response_length = 8768 + +[training] +num_epochs = 5 +learning_rate = 1e-6 +max_grad_norm = 1.0 +gradient_checkpointing = true + +[grpo] +# group_size = number of trajectories sampled per prompt. +num_samples_per_prompt = 8 +temperature = 1.0 +kl_coef = 0.0001 +clip_ratio = 0.2 +# Per-rollout tool-use turn cap (one assistant turn per loop iteration). +max_turns = 10 + +[data] +grpo_max_prompts = 20000 # pooled GRPO prompts (cached on the trainer) +prompts_per_step = 256 # prompts per _grpo_step (256 * 8 = 2048 traj/step) + +[reward] +# Cost-aware reward: R = base(+1/-1) - lam * (cost_usd / cost_max). +# `lam` is SWEPT across runs to trace the accuracy/cost Pareto frontier, +# e.g. [0.0, 0.05, 0.1, 0.2, 0.4]. The value below is the default run. +lam = 0.0 +cost_max = 0.10 + +# Optional: map local OSS model ids -> served vLLM base_url so the local-model +# tools in the orchestrator catalog dispatch to a real endpoint. Omit to leave +# them unbacked (tool still listed). +[local_endpoints] +# "Qwen/Qwen3.5-9B" = "http://localhost:8001/v1" + +[checkpoint] +checkpoint_dir = "checkpoints/orchestrator_grpo_qwen3.5_9b" +save_every_n_epochs = 1 diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml new file mode 100644 index 00000000..f3d325b0 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml @@ -0,0 +1,34 @@ +# Orchestrator: Qwen3.5-9B. v1 = base self-sampling (the served orchestrator_model is this same base). +# +# v1 data = base Qwen3.5-9B self-sampling: roll the base model out over +# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the +# dataset first (no GPU, needs the vLLM orchestrator endpoint up): +# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \ +# --out data/orchestrator_qwen3.5-9b_sft_v1.jsonl --samples-per-task 8 +# then train on a rented GPU with regenerate_traces = false. + +[model] +model_name = "Qwen/Qwen3.5-9B" +max_seq_length = 4096 + +[training] +num_epochs = 3 +batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch +learning_rate = 2e-5 +weight_decay = 0.01 +warmup_ratio = 0.1 +gradient_checkpointing = true + +[data] +trace_cache_path = "data/orchestrator_qwen3.5-9b_sft_v1.jsonl" +regenerate_traces = false +# Base Qwen3.5-9B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py). +orchestrator_endpoint = "http://localhost:8001/v1" +orchestrator_model = "qwen3.5-9b" +samples_per_task = 8 +max_keep_per_task = 1 +distill_max_tasks = 0 # 0 -> all of load_sft_tasks() + +[checkpoint] +checkpoint_dir = "checkpoints/orchestrator_qwen3.5-9b_sft" +save_every_n_epochs = 1 diff --git a/src/openjarvis/learning/intelligence/orchestrator/eval_backend.py b/src/openjarvis/learning/intelligence/orchestrator/eval_backend.py new file mode 100644 index 00000000..a7888a7e --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/eval_backend.py @@ -0,0 +1,188 @@ +"""Eval backend that drives the ToolOrchestra orchestrator as the "model". + +This adapts our unified-tool orchestrator rollout +(:func:`openjarvis.agents.hybrid.toolorchestra.rollout.run_unified_rollout`) +to the eval framework's :class:`~openjarvis.evals.core.backend.InferenceBackend` +interface so it can be scored by the existing +:class:`~openjarvis.evals.core.runner.EvalRunner` over any registered benchmark. + +The orchestrator is an OpenAI-compatible chat model (served locally via vLLM, +default ``http://localhost:8001/v1``) that emits tool calls over the fixed +8-tool catalog from :func:`expert_registry.orchestrator_catalog`. For each eval +sample we run one full rollout and return ``rollout.final_answer`` as the model +answer, plus token/cost telemetry in the ``generate_full`` payload. + +Construction is import-safe and network-free; the OpenAI SDK is only touched +inside ``call_orchestrator`` at generate time (lazy import in ``unified.py``). +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, Optional + +from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog +from openjarvis.agents.hybrid.toolorchestra import rollout as _rollout_mod +from openjarvis.agents.hybrid.toolorchestra.unified import ( + make_call_orchestrator, + make_dispatch, +) +from openjarvis.evals.core.backend import InferenceBackend + +DEFAULT_ENDPOINT = "http://localhost:8001/v1" +DEFAULT_MODEL = "qwen3-8b" + + +class OrchestratorBackend(InferenceBackend): + """Run the ToolOrchestra orchestrator rollout as an eval "model". + + Parameters + ---------- + orchestrator_endpoint: + OpenAI-compatible base URL for the served orchestrator (vLLM). Defaults + to ``http://localhost:8001/v1``. + orchestrator_model: + Served model id for the orchestrator. Defaults to ``"qwen3-8b"`` (will + be the Qwen3.5-9B checkpoint later). + api_key: + API key for the endpoint. ``"EMPTY"`` for a local vLLM server. + local_endpoints: + Optional dict mapping a local OSS model id (e.g. ``"Qwen/Qwen3.5-9B"``) + to its vLLM base URL. Unmapped local models are still listed in the + catalog but served as unconfigured (``base_url=None``). + max_turns: + Maximum orchestrator reasoning->action->observation turns per sample. + temperature: + Orchestrator sampling temperature. + dispatch_cfg: + Optional config dict passed to :func:`make_dispatch` (worker dispatch). + """ + + backend_id = "orchestrator" + framework_name = "openjarvis-orchestrator" + + def __init__( + self, + *, + orchestrator_endpoint: str = DEFAULT_ENDPOINT, + orchestrator_model: str = DEFAULT_MODEL, + api_key: str = "EMPTY", + local_endpoints: Optional[Dict[str, str]] = None, + max_turns: int = 8, + temperature: float = 1.0, + dispatch_cfg: Optional[Dict[str, Any]] = None, + ) -> None: + self.orchestrator_endpoint = orchestrator_endpoint + self.orchestrator_model = orchestrator_model + self.api_key = api_key + self.local_endpoints = dict(local_endpoints or {}) + self.max_turns = int(max_turns) + self.temperature = float(temperature) + self._dispatch_cfg = dict(dispatch_cfg or {}) + + # Build the catalog once; it is stateless. ``local_endpoints`` now maps + # full local model ids (e.g. ``"Qwen/Qwen3.5-9B"``) -> vLLM base_url. + self._tools = orchestrator_catalog(local_endpoints=self.local_endpoints) + + # ------------------------------------------------------------------ + # InferenceBackend abstract methods + # ------------------------------------------------------------------ + def generate( + self, + prompt: str, + *, + model: str, + system: str = "", + temperature: float = 0.0, + max_tokens: int = 2048, + ) -> str: + """Run a rollout and return just the final-answer text.""" + full = self.generate_full( + prompt, + model=model, + system=system, + temperature=temperature, + max_tokens=max_tokens, + ) + return full.get("content", "") or "" + + def generate_full( + self, + prompt: str, + *, + model: str, + system: str = "", + temperature: float = 0.0, + max_tokens: int = 2048, + ) -> Dict[str, Any]: + """Run one orchestrator rollout over ``prompt`` and return full details. + + Returns a dict with the keys the runner reads: ``content``, ``usage``, + ``model``, ``latency_seconds``, ``cost_usd``, plus ``tool_calls`` / + ``turn_count`` / ``framework``. On error, returns ``content=""`` and a + non-empty ``error`` field so the runner records it as a failed sample + rather than crashing the whole run. + """ + # The caller-supplied ``model`` is the RunConfig.model (the orchestrator + # model name). Fall back to the configured one if blank. + orch_model = model or self.orchestrator_model + # Prefer the explicit per-call temperature only when the caller set a + # non-default value; otherwise use the backend's configured temperature. + temp = temperature if temperature else self.temperature + + started = time.time() + try: + call_orch = make_call_orchestrator( + orch_model, + base_url=self.orchestrator_endpoint, + api_key=self.api_key, + temperature=temp, + max_tokens=max_tokens, + ) + dispatch = make_dispatch(self._dispatch_cfg) + rollout = _rollout_mod.run_unified_rollout( + prompt, + self._tools, + call_orchestrator=call_orch, + dispatch=dispatch, + max_turns=self.max_turns, + ) + except Exception as exc: # noqa: BLE001 - surface as a failed sample + return { + "content": "", + "error": f"{type(exc).__name__}: {exc}", + "usage": {"prompt_tokens": 0, "completion_tokens": 0}, + "model": orch_model, + "latency_seconds": time.time() - started, + "cost_usd": 0.0, + "tool_calls": 0, + "turn_count": 0, + "framework": self.framework_name, + } + + latency = time.time() - started + total_tokens = int(getattr(rollout, "tokens", 0) or 0) + return { + "content": rollout.final_answer or "", + # The rollout reports a single combined token count; we surface it + # as completion_tokens so it still flows into total-token telemetry. + "usage": { + "prompt_tokens": 0, + "completion_tokens": total_tokens, + }, + "model": orch_model, + "latency_seconds": latency, + "cost_usd": float(getattr(rollout, "cost_usd", 0.0) or 0.0), + "tool_calls": int(getattr(rollout, "num_tool_calls", 0) or 0), + "turn_count": len(getattr(rollout, "turns", []) or []), + "framework": self.framework_name, + "trace_data": { + "parse_failures": int(getattr(rollout, "parse_failures", 0) or 0), + "tool_calls": [ + {"name": n, "arguments": a} for n, a in rollout.tool_calls() + ], + }, + } + + +__all__ = ["OrchestratorBackend", "DEFAULT_ENDPOINT", "DEFAULT_MODEL"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py index 27e58c59..4f885be1 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py +++ b/src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py @@ -6,8 +6,23 @@ doesn't require a separate critic model — instead, it uses trajectories, compute rewards, normalise within the group, and update the policy to increase the probability of better solutions. +**Trajectory-level GRPO over the unified tool-use rollout.** Generation +flows through :func:`run_unified_rollout` (the SAME faithful +```` multi-turn loop our SFT data is built from), NOT the old +single-turn ``THOUGHT/TOOL/INPUT`` environment path. The policy *is* the +orchestrator: each turn we tokenize ``system+user`` with the model's chat +template, generate ONE assistant turn, parse any +``{...}``, and record the exact +``(input_ids, generated_ids)`` for that turn. Once a rollout finishes, the +recorded turns are stitched into one trajectory token sequence with an +*assistant mask* (1 only on policy-generated tokens) so the policy gradient +is computed solely over the tokens the model actually authored, under both +the current policy (with grad) and the frozen reference (no grad). + All ``torch``/``transformers`` imports are guarded so the module can be -imported without GPU dependencies. +imported without GPU dependencies; the rollout + log-prob paths are +monkeypatchable so ``_grpo_step``/``_train_epoch`` are CPU-smoke-testable +with fakes (``policy.model is None``). """ from __future__ import annotations @@ -17,7 +32,7 @@ import logging import shutil from dataclasses import asdict, dataclass, field from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional, Tuple # Optional imports ----------------------------------------------------------- try: @@ -68,11 +83,24 @@ class OrchestratorGRPOConfig: max_grad_norm: float = 1.0 # GRPO specific + # ``num_samples_per_prompt`` is the GRPO *group size*: N trajectories + # sampled per prompt for group-relative advantages. num_samples_per_prompt: int = 8 temperature: float = 1.0 kl_coef: float = 0.0001 clip_ratio: float = 0.2 + # Data / batching + # How many prompts to pool from the GRPO dataset (cached on the trainer). + grpo_max_prompts: int = 20000 + # How many prompts go into one ``_grpo_step`` (a batch / step). With + # group_size N this is prompts_per_step * N trajectories per step. + prompts_per_step: int = 256 + + # Cost-aware reward (r=+1/-1, R -= lam * cost / cost_max). ``lam`` is swept. + lam: float = 0.0 + cost_max: float = 0.10 + # Environment available_tools: List[str] = field( default_factory=lambda: [ @@ -83,6 +111,11 @@ class OrchestratorGRPOConfig: ] ) max_turns: int = 10 + # Optional local OSS model endpoints, mapping full model id (e.g. + # "Qwen/Qwen3.5-9B") -> vLLM base_url. Passed to ``orchestrator_catalog`` + # so the local-model tools dispatch to a served endpoint; omitted when + # unset (tool still listed, just unbacked). + local_endpoints: Dict[str, str] = field(default_factory=dict) # Checkpoint checkpoint_dir: str = "checkpoints/orchestrator_grpo" @@ -95,6 +128,24 @@ class OrchestratorGRPOConfig: use_8bit_optimizer: bool = False +@dataclass +class RolloutRecord: + """One sampled unified-rollout trajectory, with the tokens + assistant mask + needed to compute its policy log-prob. + + ``token_ids`` is the full stitched trajectory (prompt / observation tokens + interleaved with policy-generated tokens over turns). ``assistant_mask`` + is parallel: 1 only on policy-generated positions. ``advantage`` is filled + in after group-relative normalisation. + """ + + final_answer: str + cost_usd: float + token_ids: List[int] = field(default_factory=list) + assistant_mask: List[int] = field(default_factory=list) + advantage: float = 0.0 + + # --------------------------------------------------------------------------- # Trainer # --------------------------------------------------------------------------- @@ -111,12 +162,37 @@ class OrchestratorGRPOTrainer: self.device = None self.global_step = 0 + # Cost-aware reward (r=+1/-1, R -= lam * cost / cost_max). ``lam`` + # is swept across runs to trace the accuracy/cost Pareto frontier. + from openjarvis.learning.intelligence.orchestrator.reward import ( + CostAwareReward, + ) + + self._reward = CostAwareReward(lam=config.lam, cost_max=config.cost_max) + + # Lazily-loaded, cached GRPO prompt pool (list[Task]). + self._prompts: List[Any] | None = None + # Lazily-built, cached unified tool catalog (list[ExpertTool]). + self._tool_catalog: Optional[List[Any]] = None + if HAS_TORCH and torch is not None: self.device = _select_torch_device() self._init_model() self._init_optimizer() + # -- Data ---------------------------------------------------------------- + + def _load_prompts(self) -> List[Any]: + """Load (and cache) the GRPO prompt pool — a ``list[Task]``.""" + if self._prompts is None: + from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import ( # noqa: E501 + load_grpo_prompts, + ) + + self._prompts = load_grpo_prompts(n=self.config.grpo_max_prompts) + return self._prompts + # -- Initialisation ------------------------------------------------------ def _init_model(self) -> None: @@ -188,82 +264,270 @@ class OrchestratorGRPOTrainer: self._save_checkpoint(epoch) def _train_epoch(self, epoch: int) -> Dict[str, float]: - if self.policy.model is None: - return {"epoch": epoch, "loss": 0.0, "reward": 0.0} + """Iterate the GRPO prompt pool in chunks and run ``_grpo_step``. + + Prompts are loaded once via :meth:`_load_prompts` (cached on the + trainer) and chunked into batches of ``config.prompts_per_step``. + Each chunk is one ``_grpo_step`` call (group_size trajectories per + prompt). Real model work happens inside ``_grpo_step`` and is + guarded there; ``_train_epoch`` itself does no torch work, so it + can be smoke-tested on CPU by monkeypatching ``_grpo_step``. + """ + if HAS_TORCH and self.policy.model is not None: + self.policy.model.train() + + prompts = self._load_prompts() + step_size = max(int(self.config.prompts_per_step), 1) - self.policy.model.train() total_loss = 0.0 total_reward = 0.0 + total_accuracy = 0.0 + n_prompts = 0 num_batches = 0 - # In a real implementation, iterate over a task dataset. - # Here we provide the skeleton; actual data loading is - # trainer-specific. - self.global_step += 1 - num_batches = max(num_batches, 1) + for start in range(0, len(prompts), step_size): + chunk = prompts[start : start + step_size] + if not chunk: + continue + metrics = self._grpo_step(chunk) + loss = float(metrics.get("loss", 0.0)) + reward = float(metrics.get("reward", 0.0)) + accuracy = float(metrics.get("accuracy", 0.0)) + n = int(metrics.get("n_prompts", len(chunk))) - avg_loss = total_loss / num_batches if num_batches > 0 else 0.0 - avg_reward = total_reward / num_batches if num_batches > 0 else 0.0 - return {"epoch": epoch, "loss": avg_loss, "reward": avg_reward} + total_loss += loss * n + total_reward += reward * n + total_accuracy += accuracy * n + n_prompts += n + num_batches += 1 + self.global_step += 1 - def _grpo_step( - self, - prompts: List[str], - ground_truths: List[str], - ) -> tuple: - """Perform one GRPO training step. + denom = n_prompts if n_prompts > 0 else 1 + return { + "epoch": epoch, + "loss": total_loss / denom, + "reward": total_reward / denom, + "accuracy": total_accuracy / denom, + "n_prompts": n_prompts, + "n_batches": num_batches, + } - For each prompt: - 1. Generate N candidate trajectories. - 2. Compute reward for each. - 3. Normalise advantages within the group. - 4. Compute clipped policy gradient + KL penalty. + # -- Tool catalog -------------------------------------------------------- + + def _tools(self) -> List[Any]: + """The orchestrator's unified tool catalog (cached). + + Two model classes (cloud frontier + local OSS) plus the basic tools, + exactly what the rollout / SFT data condition on. ``local_endpoints`` + wires the local-model tools to their served vLLM endpoints. + """ + if getattr(self, "_tool_catalog", None) is None: + from openjarvis.agents.hybrid.expert_registry import ( + orchestrator_catalog, + ) + + self._tool_catalog = orchestrator_catalog( + local_endpoints=dict(self.config.local_endpoints) or None, + ) + return self._tool_catalog + + # -- Rollout / trajectory collection ------------------------------------- + + def _rollout_group(self, task: Any) -> List["RolloutRecord"]: + """Sample ``group_size`` rollouts for one task through the unified loop. + + Each rollout is driven by a *policy* ``call_orchestrator`` (the HF + model being trained generating one assistant turn at a time) and the + real :func:`make_dispatch` tool executor. Per turn we capture the exact + ``(input_ids, generated_ids)`` so the trajectory token sequence + + assistant mask can be rebuilt for the log-prob pass. + + Returns one :class:`RolloutRecord` per sampled rollout. Overridable / + monkeypatchable in tests so ``_grpo_step`` runs with fakes when + ``policy.model is None``. + """ + from openjarvis.agents.hybrid.toolorchestra.rollout import ( + build_system_prompt, + run_unified_rollout, + ) + from openjarvis.agents.hybrid.toolorchestra.unified import make_dispatch + from openjarvis.agents.hybrid.expert_registry import build_tool_specs + + tools = self._tools() + system = build_system_prompt(build_tool_specs(tools)) + dispatch = make_dispatch({}) + + records: List[RolloutRecord] = [] + for _ in range(self.config.num_samples_per_prompt): + turn_buffer: List[Tuple[List[int], List[int]]] = [] + call_orchestrator = self._make_policy_caller(turn_buffer) + rollout = run_unified_rollout( + task.instruction, + tools, + call_orchestrator=call_orchestrator, + dispatch=dispatch, + max_turns=self.config.max_turns, + system=system, + ) + token_ids, assistant_mask = self._stitch_trajectory(turn_buffer) + records.append( + RolloutRecord( + final_answer=rollout.final_answer, + cost_usd=float(rollout.cost_usd), + token_ids=token_ids, + assistant_mask=assistant_mask, + ) + ) + return records + + def _make_policy_caller(self, turn_buffer): + """Build a policy-driven ``call_orchestrator(system, user, specs)``. + + Tokenizes ``system+user`` with the model's chat template, generates ONE + assistant turn with the current policy, parses any + ``{...}``, appends the per-turn + ``(input_ids, generated_ids)`` to ``turn_buffer`` (for the later + log-prob pass), and returns + ``(text, tool_calls, prompt_tokens, completion_tokens)`` as the unified + rollout expects. + """ + from openjarvis.agents.hybrid.toolorchestra.parsing import ( + _parse_rl_tool_call, + ) + + tok = self.policy.tokenizer + + def call_orchestrator(system: str, user: str, specs): + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ] + input_ids = tok.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True + ) + inputs = { + "input_ids": torch.tensor([input_ids], device=self.device), + } + input_len = len(input_ids) + max_new = min(self.config.max_response_length, 32000 - input_len - 100) + max_new = max(min(max_new, 2048), 16) + + with torch.no_grad(): + out = self.policy.model.generate( + **inputs, + max_new_tokens=max_new, + temperature=self.config.temperature, + do_sample=True, + ) + gen_ids = out[0][input_len:].tolist() + text = tok.decode(gen_ids, skip_special_tokens=True) + + # Record exact tokens for the trajectory log-prob pass. + turn_buffer.append((list(input_ids), list(gen_ids))) + + parsed = _parse_rl_tool_call(text, None) + tool_calls = ( + [(parsed["name"], parsed["arguments"])] if parsed else [] + ) + return text, tool_calls, int(input_len), int(len(gen_ids)) + + return call_orchestrator + + def _stitch_trajectory( + self, turn_buffer: List[Tuple[List[int], List[int]]] + ) -> Tuple[List[int], List[int]]: + """Concatenate per-turn tokens into one trajectory + assistant mask. + + For turn ``i`` the buffer holds ``(prompt_ids_i, gen_ids_i)`` where + ``prompt_ids_i`` already contains all prior turns' assistant text and + the tool observations re-rendered into the chat by the rollout loop + (the rollout rebuilds ``user`` each turn from the running context). So + the faithful trajectory token stream is, for each turn, the *new* + prefix tokens (mask 0 — prompt / observation / chat-template tokens) + followed by the generated assistant tokens (mask 1). + + We reconstruct it by diffing successive prompts: turn ``i`` contributes + the prompt tokens that weren't already emitted, then its generated + tokens. This keeps the assistant mask exact (1 only on policy-generated + positions) without re-tokenizing the whole transcript. + """ + token_ids: List[int] = [] + mask: List[int] = [] + emitted = 0 # length of token_ids already produced + for prompt_ids, gen_ids in turn_buffer: + # New prompt-side tokens for this turn (everything past what we've + # already emitted). These are observation / chat-template tokens → + # mask 0 (not policy-authored). + new_prompt = prompt_ids[emitted:] if len(prompt_ids) > emitted else [] + token_ids.extend(new_prompt) + mask.extend([0] * len(new_prompt)) + # Generated assistant tokens → mask 1. + token_ids.extend(gen_ids) + mask.extend([1] * len(gen_ids)) + emitted = len(token_ids) + return token_ids, mask + + def _grpo_step(self, tasks: List[Any]) -> Dict[str, float]: + """Perform one GRPO training step over a batch of ``Task`` objects. + + For each task: + 1. Sample N candidate trajectories (``group_size`` = + ``num_samples_per_prompt``) through the unified tool-use rollout — + the SAME ```` multi-turn path the SFT data uses, driven + by the *policy* as orchestrator (see :meth:`_rollout_group`). + 2. Verify each rollout's ``final_answer`` against the gold answer + (:func:`verify_answer`) and score with :class:`CostAwareReward` + over an :class:`Episode` carrying ``.correct`` and + ``.total_cost_usd``. + 3. Normalise advantages within the group (mean/std). + 4. Compute the trajectory-level clipped policy gradient + KL penalty. + The trajectory log-prob is the sum of per-token log-probs over the + *assistant-masked* positions (:meth:`_trajectory_logprob`). 5. Backward + clip + step. - Returns ``(loss_value, avg_reward)``. + Returns a metrics dict with ``loss``, ``reward``, ``accuracy``, + and ``n_prompts``. """ if self.policy.model is None or not HAS_TORCH: raise RuntimeError("Cannot train without PyTorch and model.") self.policy.model.train() - all_prompts: list[str] = [] - all_responses: list[str] = [] - all_advantages: list[float] = [] - all_rewards: list[float] = [] - - from openjarvis.learning.intelligence.orchestrator.reward import ( - MultiObjectiveReward, - Normalizers, - RewardWeights, + from openjarvis.learning.intelligence.orchestrator.sft_data.verify import ( + verify_answer, ) from openjarvis.learning.intelligence.orchestrator.types import ( Episode, ) - reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers()) + reward_fn = self._reward - for prompt, gt in zip(prompts, ground_truths): + # (token_ids, assistant_mask, advantage) per trajectory, plus running + # reward / accuracy stats. + traj_records: list[RolloutRecord] = [] + all_advantages: list[float] = [] + all_rewards: list[float] = [] + all_correct: list[float] = [] + + for task in tasks: + records = self._rollout_group(task) group_rewards: list[float] = [] - group_responses: list[str] = [] - for _ in range(self.config.num_samples_per_prompt): - response, _log_probs = self._generate_with_log_probs(prompt) - - # Build a minimal episode for reward computation + for rec in records: + correct = verify_answer(task, rec.final_answer) episode = Episode( - task_id="grpo", - initial_prompt=prompt, - ground_truth=gt, - final_answer=response, - correct=(response.strip() == gt.strip()), + task_id=getattr(task, "task_id", ""), + initial_prompt=task.instruction, + ground_truth=getattr(task, "answer", ""), + final_answer=rec.final_answer, + correct=bool(correct), + total_cost_usd=rec.cost_usd, ) reward = reward_fn.compute(episode) - group_rewards.append(reward) - group_responses.append(response) + all_correct.append(1.0 if correct else 0.0) - # Group-relative advantages + # Group-relative advantages (mean/std within the group). mean_r = sum(group_rewards) / len(group_rewards) std_r = ( sum((r - mean_r) ** 2 for r in group_rewards) / len(group_rewards) @@ -273,21 +537,25 @@ class OrchestratorGRPOTrainer: else: advantages = [0.0] * len(group_rewards) - for resp, adv, rew in zip(group_responses, advantages, group_rewards): - all_prompts.append(prompt) - all_responses.append(resp) + for rec, adv, rew in zip(records, advantages, group_rewards): + rec.advantage = adv + traj_records.append(rec) all_advantages.append(adv) all_rewards.append(rew) - # Policy gradient loss + # Trajectory-level policy gradient loss. total_loss = torch.tensor(0.0, device=self.device, requires_grad=True) - for prompt, response, advantage in zip( - all_prompts, all_responses, all_advantages - ): - current_lp = self._compute_log_probs(prompt, response) + for rec in traj_records: + if not rec.token_ids or sum(rec.assistant_mask) == 0: + continue + current_lp = self._trajectory_logprob( + rec.token_ids, rec.assistant_mask, ref=False + ) with torch.no_grad(): - ref_lp = self._compute_log_probs_ref(prompt, response) + ref_lp = self._trajectory_logprob( + rec.token_ids, rec.assistant_mask, ref=True + ) log_ratio = current_lp - ref_lp ratio = torch.exp(log_ratio) @@ -296,16 +564,25 @@ class OrchestratorGRPOTrainer: clip = self.config.clip_ratio clipped = torch.clamp(ratio, 1 - clip, 1 + clip) - policy_loss = -torch.min(ratio * advantage, clipped * advantage) + policy_loss = -torch.min(ratio * rec.advantage, clipped * rec.advantage) kl = (ratio - 1) - log_ratio total_loss = total_loss + policy_loss + self.config.kl_coef * kl - avg_loss = total_loss / max(len(all_prompts), 1) + n_prompts = len(tasks) + avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0 + avg_acc = sum(all_correct) / len(all_correct) if all_correct else 0.0 + + n_traj = max(len(traj_records), 1) + avg_loss = total_loss / n_traj loss_val = avg_loss.item() if torch.isnan(avg_loss) or torch.isinf(avg_loss): - avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0 - return 0.0, float(avg_reward) + return { + "loss": 0.0, + "reward": float(avg_reward), + "accuracy": float(avg_acc), + "n_prompts": n_prompts, + } self.optimizer.zero_grad() avg_loss.backward() @@ -314,107 +591,68 @@ class OrchestratorGRPOTrainer: for param in self.policy.model.parameters(): if param.grad is not None and torch.isnan(param.grad).any(): self.optimizer.zero_grad() - avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0 - return float(loss_val), float(avg_reward) + return { + "loss": float(loss_val), + "reward": float(avg_reward), + "accuracy": float(avg_acc), + "n_prompts": n_prompts, + } torch.nn.utils.clip_grad_norm_( self.policy.model.parameters(), self.config.max_grad_norm ) self.optimizer.step() - avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0 - return float(loss_val), float(avg_reward) + return { + "loss": float(loss_val), + "reward": float(avg_reward), + "accuracy": float(avg_acc), + "n_prompts": n_prompts, + } - # -- Generation / log-prob helpers --------------------------------------- + # -- Trajectory log-prob ------------------------------------------------- - def _generate_with_log_probs(self, prompt: str) -> "tuple[str, Any]": - """Generate a response and return ``(text, log_probs)``.""" - inputs = self.policy.tokenizer( - prompt, - return_tensors="pt", - truncation=True, - max_length=min(self.config.max_prompt_length, 16000), - ).to(self.device) + def _trajectory_logprob( + self, + token_ids: List[int], + assistant_mask: List[int], + *, + ref: bool, + ) -> "torch.Tensor": + """Sum of per-token log-probs over the assistant-masked positions. - input_len = inputs.input_ids.shape[1] - max_new = min(self.config.max_response_length, 32000 - input_len - 100) - max_new = max(min(max_new, 2048), 128) + ``token_ids`` is the full trajectory (prompt + observation + assistant + tokens interleaved over turns); ``assistant_mask[i] == 1`` iff token + ``i`` was generated by the policy. We run the model once over the whole + sequence and accumulate ``log p(token_{i} | token_{ "torch.Tensor": - """Log-probs of *response* given *prompt* under current policy.""" - full = prompt + response - inputs = self.policy.tokenizer(full, return_tensors="pt", truncation=True).to( - self.device - ) - prompt_inputs = self.policy.tokenizer( - prompt, return_tensors="pt", truncation=True - ).to(self.device) - - with torch.enable_grad(): - logits = self.policy.model(**inputs).logits - - start = len(prompt_inputs.input_ids[0]) - end = len(inputs.input_ids[0]) + if ref: + ctx = torch.no_grad() + else: + ctx = torch.enable_grad() + with ctx: + logits = model(input_ids=ids).logits lps = [] - for i in range(start, end - 1): - lp = F.log_softmax(logits[0, i, :], dim=-1)[inputs.input_ids[0, i + 1]] + # Predict token i from logits at position i-1; include it iff token i is + # an assistant-generated token. + for i in range(1, len(token_ids)): + if assistant_mask[i] != 1: + continue + lp = F.log_softmax(logits[0, i - 1, :], dim=-1)[token_ids[i]] lps.append(lp) - return torch.stack(lps).sum() if lps else torch.tensor(0.0) - - def _compute_log_probs_ref(self, prompt: str, response: str) -> "torch.Tensor": - """Log-probs under the frozen reference policy (no grad).""" - full = prompt + response - inputs = self.ref_policy.tokenizer( - full, return_tensors="pt", truncation=True - ).to(self.device) - prompt_inputs = self.ref_policy.tokenizer( - prompt, return_tensors="pt", truncation=True - ).to(self.device) - - with torch.no_grad(): - logits = self.ref_policy.model(**inputs).logits - - start = len(prompt_inputs.input_ids[0]) - end = len(inputs.input_ids[0]) - - lps = [] - for i in range(start, end - 1): - lp = F.log_softmax(logits[0, i, :], dim=-1)[inputs.input_ids[0, i + 1]] - lps.append(lp) - - return torch.stack(lps).sum() if lps else torch.tensor(0.0) + if not lps: + return torch.tensor(0.0, device=self.device) + return torch.stack(lps).sum() # -- Checkpointing ------------------------------------------------------- @@ -488,4 +726,5 @@ _ensure_registered() __all__ = [ "OrchestratorGRPOConfig", "OrchestratorGRPOTrainer", + "RolloutRecord", ] diff --git a/src/openjarvis/learning/intelligence/orchestrator/reward.py b/src/openjarvis/learning/intelligence/orchestrator/reward.py index 7cf41de5..2d756980 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/reward.py +++ b/src/openjarvis/learning/intelligence/orchestrator/reward.py @@ -143,6 +143,61 @@ class MultiObjectiveReward: return [self.compute(ep) for ep in episodes] +@dataclass +class CostAwareReward: + """Cost-aware GRPO reward: ``r = +1/-1`` for correctness, minus a + cost penalty. + + This is the headline objective from the plan:: + + base = +1.0 if episode.correct else -1.0 + reward = base - lambda * (total_cost_usd / cost_max) + + ``lambda`` (``lam``) is the cost-penalty weight and is *swept* across a + set of values (see :func:`lambda_sweep`) to trace the accuracy/cost + Pareto frontier. ``cost_max`` normalises USD cost so the penalty is + O(1) when an episode hits the per-task cost budget. + """ + + lam: float = 0.0 + """Cost-penalty weight (swept, e.g. [0.0, 0.05, 0.1, 0.2, 0.4]).""" + + cost_max: float = 0.10 + """USD cost normalizer (per-task cost budget).""" + + def compute(self, episode: Episode) -> float: + """Scalar reward: ``+1/-1`` for correctness minus a cost term.""" + base = 1.0 if episode.correct else -1.0 + return base - self.lam * (episode.total_cost_usd / self.cost_max) + + def compute_with_breakdown(self, episode: Episode) -> Dict[str, float]: + """Reward with per-component breakdown.""" + correct = 1.0 if episode.correct else 0.0 + base = 1.0 if episode.correct else -1.0 + cost_term = -self.lam * (episode.total_cost_usd / self.cost_max) + return { + "correct": correct, + "base": base, + "cost_term": cost_term, + "reward": base + cost_term, + } + + def compute_batch(self, episodes: List[Episode]) -> List[float]: + """Compute rewards for a batch of episodes.""" + return [self.compute(ep) for ep in episodes] + + +def lambda_sweep( + values: List[float], cost_max: float = 0.10 +) -> List["CostAwareReward"]: + """Build one :class:`CostAwareReward` per lambda value. + + Used to sweep the cost-penalty weight and trace the accuracy/cost + Pareto frontier (the ``R -= lambda * cost / cost_max`` objective). + """ + return [CostAwareReward(lam=lam, cost_max=cost_max) for lam in values] + + class AdaptiveRewardWeights: """Adaptive reward weights that shift during training. @@ -211,7 +266,9 @@ class AdaptiveRewardWeights: __all__ = [ "AdaptiveRewardWeights", + "CostAwareReward", "MultiObjectiveReward", "Normalizers", "RewardWeights", + "lambda_sweep", ] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py index 5e7f5338..253eea40 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py @@ -1,55 +1,44 @@ -"""Synthetic SFT-data generation for the orchestrator cold-start. +"""SFT-data generation for the orchestrator cold-start. -Turns NeuLab ADP (``neulab/agent-data-collection``) agent trajectories into -THOUGHT/TOOL/INPUT ``conversations`` JSONL that +Execution-grounded rejection sampling in the faithful ToolOrchestra action +space (arXiv:2511.21689): for each ``nvidia/ToolScale`` task, roll out a teacher +orchestrator over the unified tool catalog N times, verify each trajectory, keep +the cheapest passing one(s), and serialize them into the ```` +``conversations`` JSONL that :class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset` loads directly. Pipeline:: - ADP trajectory -> canonical Episode -> per-paradigm tiered renderings - -> reward-ranked best-correct rendering -> conversations JSONL + ToolScale task -> N teacher rollouts -> verify -> keep cheapest passing + -> conversations JSONL -Everything here runs with **no GPU and no API keys** (the cold-start does not -re-execute models; it re-tiers the demonstrated ADP traces). +(The earlier ADP-relabel cold-start was a heuristic re-tiering of demonstrated +traces; it was removed in favour of this grounded pipeline.) """ from __future__ import annotations -from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import ( - CanonicalStep, - iter_trajectories, - trajectory_rows_to_episode, +from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import ( + generate_sft_dataset, + gold_coverage_verify, ) -from openjarvis.learning.intelligence.orchestrator.sft_data.build import ( - build_sft_dataset, +from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import ( + GoldAction, + ToolScaleTask, + load_toolscale, + normalize_row, ) -from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import ( - PARADIGMS, - RenderedEpisode, - render_all, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best -from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record -from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import ( - Difficulty, - Tier, - step_difficulty, - tier_telemetry, +from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import ( + trajectory_to_record, ) __all__ = [ - "CanonicalStep", - "Difficulty", - "PARADIGMS", - "RenderedEpisode", - "Tier", - "build_sft_dataset", - "iter_trajectories", - "render_all", - "select_best", - "step_difficulty", - "tier_telemetry", - "to_record", - "trajectory_rows_to_episode", + "GoldAction", + "ToolScaleTask", + "generate_sft_dataset", + "gold_coverage_verify", + "load_toolscale", + "normalize_row", + "trajectory_to_record", ] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py deleted file mode 100644 index b37e5313..00000000 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py +++ /dev/null @@ -1,198 +0,0 @@ -"""Full-trajectory reader for NeuLab ADP (``neulab/agent-data-collection``). - -The eval-side loader (:mod:`openjarvis.evals.datasets.adp`) keeps only -``(problem, reference)`` per row — it drops every intermediate step, which is -exactly the signal the orchestrator needs. This module instead transcribes -**all** turns of a trajectory into a canonical -:class:`~openjarvis.learning.intelligence.orchestrator.types.Episode`. - -No GPU / no API keys: we read the demonstrated traces and re-tier them -downstream; we never re-execute a model here. -""" - -from __future__ import annotations - -import ast -from dataclasses import dataclass -from typing import Iterable, Iterator, List, MutableMapping, Optional - -from openjarvis.learning.intelligence.orchestrator.types import ( - Episode, - OrchestratorAction, - OrchestratorObservation, -) - -HF_DATASET_ID = "neulab/agent-data-collection" -HF_SPLIT = "std" - -# Same configs the eval loader concatenates; the meeting steer favours the -# coding / agentic / tool-use ones. -DEFAULT_CONFIGS: tuple[str, ...] = ( - "codeactinstruct", - "code_feedback", - "openhands", - "agenttuning_os", - "agenttuning_db", - "swe-smith", -) - -# ADP turn ``class_`` values that denote a tool/code action vs. plain message. -_CODE_CLASSES = {"code_action", "ipython_action", "bash_action"} -_SEARCH_CLASSES = {"search_action", "browse_action", "web_action"} -_MESSAGE_CLASSES = {"message_action"} - - -@dataclass -class CanonicalStep: - """One transcribed ADP turn: what the demonstration did at this step.""" - - kind: str - """Normalised step kind: ``reason`` | ``code`` | ``search`` | ``message``.""" - - content: str - """The turn's text (the action the demonstration took).""" - - observation: str = "" - """The environment/tool result that followed, if any.""" - - is_final: bool = False - """Whether this is the trajectory's final answer turn.""" - - -def _parse_content(raw: object) -> List[MutableMapping[str, object]]: - """Parse the ``content`` field (a list, or a string repr of one).""" - if isinstance(raw, list): - return raw # type: ignore[return-value] - if isinstance(raw, str): - try: - parsed = ast.literal_eval(raw) - if isinstance(parsed, list): - return parsed # type: ignore[return-value] - except (ValueError, SyntaxError): - pass - return [] - - -def _classify(turn: MutableMapping[str, object]) -> str: - cls = str(turn.get("class_") or "").lower() - if cls in _CODE_CLASSES: - return "code" - if cls in _SEARCH_CLASSES: - return "search" - if cls in _MESSAGE_CLASSES: - return "message" - return "reason" - - -def trajectory_rows_to_episode( - record_id: str, - turns: List[MutableMapping[str, object]], -) -> Optional[Episode]: - """Transcribe one ADP trajectory's turns into a canonical ``Episode``. - - The orchestrator-relevant structure is preserved: the first user turn is the - problem; every subsequent agent turn becomes a step (with its following - observation, if the trace recorded one). Returns ``None`` if there is no - usable problem or no agent steps. - """ - problem: Optional[str] = None - steps: List[CanonicalStep] = [] - - # Group agent turns with the observation/user turn that follows them. - pending: Optional[CanonicalStep] = None - for turn in turns: - source = str(turn.get("source") or "").lower() - text = str(turn.get("content") or "").strip() - if not text: - continue - - if source == "user": - if problem is None: - problem = text - elif pending is not None: - # A user/environment turn after an agent action = its observation. - pending.observation = text[:2000] - continue - - # Agent-sourced turn -> a new step. - if pending is not None: - steps.append(pending) - pending = CanonicalStep(kind=_classify(turn), content=text) - - if pending is not None: - steps.append(pending) - - if not problem or not steps: - return None - - steps[-1].is_final = True - steps[-1].kind = "message" if steps[-1].kind == "reason" else steps[-1].kind - - episode = Episode( - task_id=record_id, - initial_prompt=problem, - ground_truth=steps[-1].content[:2000], - # ADP rows are demonstrated solutions; treat them as correct teachers. - correct=True, - ) - for st in steps: - episode.add_step( - OrchestratorAction( - thought="", # filled in by the paradigm renderer - tool_name=st.kind, - tool_input=st.content, - is_final_answer=st.is_final, - ), - OrchestratorObservation(content=st.observation or st.content), - ) - episode.final_answer = steps[-1].content[:2000] - episode.metadata["source"] = "adp" - return episode - - -def iter_trajectories( - *, - max_tasks: Optional[int] = None, - configs: Iterable[str] = DEFAULT_CONFIGS, - min_steps: int = 1, - max_steps: int = 24, -) -> Iterator[Episode]: - """Stream canonical ``Episode`` objects from ADP. - - Network is touched lazily inside the loop (``datasets.load_dataset`` with - ``streaming=True``) so importing this module stays free. Configs that fail - to load (gated/missing) are skipped. - """ - from datasets import load_dataset - - emitted = 0 - for cfg in configs: - if max_tasks is not None and emitted >= max_tasks: - break - try: - stream = load_dataset(HF_DATASET_ID, cfg, split=HF_SPLIT, streaming=True) - except Exception: - continue - for i, row in enumerate(stream): - if max_tasks is not None and emitted >= max_tasks: - break - row = dict(row) # type: ignore[arg-type] - turns = _parse_content(row.get("content")) - row_id = row.get("id") - rec_id = str(row_id) if row_id is not None else f"{cfg}-{i}" - episode = trajectory_rows_to_episode(rec_id, turns) - if episode is None: - continue - n = episode.num_turns() - if n < min_steps or n > max_steps: - continue - emitted += 1 - yield episode - - -__all__ = [ - "CanonicalStep", - "DEFAULT_CONFIGS", - "iter_trajectories", - "trajectory_rows_to_episode", -] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py deleted file mode 100644 index fe213be9..00000000 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py +++ /dev/null @@ -1,112 +0,0 @@ -"""End-to-end SFT dataset builder + CLI. - - ADP trajectory -> canonical Episode -> render_all -> select_best - -> to_record -> JSONL (+ a sidecar ``.stats.json``) - -Runs with no GPU and no API keys (the cold-start re-tiers demonstrated traces; -it does not execute models). Network is only touched to stream ADP rows. -""" - -from __future__ import annotations - -import argparse -import json -import logging -from collections import Counter -from pathlib import Path -from typing import Callable, Iterable, Iterator, Optional - -from openjarvis.learning.intelligence.orchestrator.reward import ( - MultiObjectiveReward, - Normalizers, - RewardWeights, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import ( - DEFAULT_CONFIGS, - iter_trajectories, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import render_all -from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best -from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record -from openjarvis.learning.intelligence.orchestrator.types import Episode - -logger = logging.getLogger(__name__) - - -def build_sft_dataset( - out_path: str, - *, - max_tasks: Optional[int] = 2000, - configs: Iterable[str] = DEFAULT_CONFIGS, - source: Optional[Callable[..., Iterator[Episode]]] = None, -) -> dict: - """Build the SFT JSONL at ``out_path`` and return stats. - - ``source`` overrides the ADP stream (used by tests to inject fixtures); it - must be a callable returning an iterator of canonical :class:`Episode`. - """ - reward = MultiObjectiveReward(RewardWeights(), Normalizers()) - out = Path(out_path) - out.parent.mkdir(parents=True, exist_ok=True) - - episodes = ( - source(max_tasks=max_tasks, configs=configs) - if source is not None - else iter_trajectories(max_tasks=max_tasks, configs=configs) - ) - - seen = 0 - written = 0 - dropped = 0 - paradigm_counts: Counter[str] = Counter() - - with out.open("w") as fh: - for episode in episodes: - seen += 1 - best = select_best(render_all(episode), reward=reward) - if best is None: - dropped += 1 - continue - record = to_record(best, reward=reward.compute(best.episode)) - fh.write(json.dumps(record) + "\n") - written += 1 - paradigm_counts[best.paradigm] += 1 - - stats = { - "out_path": str(out), - "tasks_seen": seen, - "records_written": written, - "tasks_dropped": dropped, - "paradigm_distribution": dict(paradigm_counts), - } - stats_path = out.with_suffix(out.suffix + ".stats.json") - stats_path.write_text(json.dumps(stats, indent=2)) - logger.info("Wrote %d SFT records to %s (%s)", written, out, dict(paradigm_counts)) - return stats - - -def _main(argv: Optional[list[str]] = None) -> int: - parser = argparse.ArgumentParser( - description="Build orchestrator SFT data from ADP." - ) - parser.add_argument("--out", default="data/orchestrator_sft_traces.jsonl") - parser.add_argument("--max-tasks", type=int, default=2000) - parser.add_argument( - "--adp-configs", - default=",".join(DEFAULT_CONFIGS), - help="comma-separated ADP sub-configs to stream", - ) - args = parser.parse_args(argv) - - logging.basicConfig(level=logging.INFO, format="%(message)s") - configs = [c.strip() for c in args.adp_configs.split(",") if c.strip()] - stats = build_sft_dataset(args.out, max_tasks=args.max_tasks, configs=configs) - print(json.dumps(stats, indent=2)) - return 0 - - -if __name__ == "__main__": # pragma: no cover - raise SystemExit(_main()) - - -__all__ = ["build_sft_dataset"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py new file mode 100644 index 00000000..4a9193d1 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py @@ -0,0 +1,282 @@ +"""Loaders for the orchestrator reasoning-SFT task sources. + +Two HuggingFace reasoning datasets give us verifiable question/answer tasks for +the small orchestrator's cold-start SFT and the GRPO prompt pool: + +- ``natolambert/GeneralThought-430K-filtered`` — open reasoning traces scraped + from gr.inc. Each row carries a ``question``, a short ``reference_answer`` + (the gold), a long ``model_answer`` (R1's full solution), and provenance in + ``question_source`` / ``task``. There is **no** clean single "category" field; + we derive a coarse ``domain`` (math / medical / code / chat / misc) from + ``question_source`` so the verifier can pick the right checker. + +- ``open-thoughts/OpenThoughts3-1.2M`` — OpenThoughts3 distillation set. Each + row has a ``domain`` in {code, math, science}, a ``source``, a ``difficulty``, + and a ``conversations`` list ``[{from: human, value}, {from: gpt, value}]``. + The human turn is the question; the gpt turn is a ``…`` trace + followed by the final solution, from which we extract the gold answer. + +Mirrors the style of the sibling ``hotpotqa.py`` / ``toolscale.py`` loaders: a +plain dataclass plus loaders that accept a ``source=`` iterable override so the +normalization path is exercised offline with no network. Network imports of +``datasets`` are kept lazy (inside the function). +""" + +from __future__ import annotations + +import random +import re +from dataclasses import dataclass +from typing import Any, Dict, Iterable, Iterator, List, Optional + +GENERALTHOUGHT_ID = "natolambert/GeneralThought-430K-filtered" +OPENTHOUGHTS_ID = "open-thoughts/OpenThoughts3-1.2M" + + +@dataclass +class Task: + task_id: str + question: str + answer: str + domain: str + + @property + def instruction(self) -> str: + return self.question + + +# --- GeneralThought ----------------------------------------------------------- + +# question_source -> coarse domain. GeneralThought has no literal category field; +# the source string is the most reliable signal (NuminaMath is math, NHSQA is +# medical, TACO/glaive are code, oasst1/lmsys are open chat, everything else misc). +_SOURCE_DOMAIN = ( + ("numina", "math"), + ("math", "math"), + ("nhsqa", "medical"), + ("medical", "medical"), + ("medicine", "medical"), + ("taco", "code"), + ("code", "code"), + ("glaive", "code"), + ("oasst", "chat"), + ("lmsys", "chat"), + ("chat", "chat"), +) + + +def _generalthought_domain(row: Dict[str, Any]) -> str: + src = str(row.get("question_source") or "").lower() + task = str(row.get("task") or "").lower() + hay = f"{src} {task}" + for needle, dom in _SOURCE_DOMAIN: + if needle in hay: + return dom + return "misc" + + +def _normalize_generalthought(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]: + question = str(row.get("question") or "").strip() + # Prefer the short reference answer (exact-match-able); fall back to the + # long model_answer when no reference is present. + answer = str(row.get("reference_answer") or "").strip() + if not answer: + answer = str(row.get("model_answer") or "").strip() + if not question or not answer: + return None + task_id = str(row.get("question_id") or f"generalthought-{index}") + return Task( + task_id=task_id, + question=question, + answer=answer, + domain=_generalthought_domain(row), + ) + + +def load_generalthought( + *, + n: int = 2000, + seed: int = 42, + source: Optional[Iterable[Dict[str, Any]]] = None, +) -> Iterator[Task]: + """Yield up to ``n`` GeneralThought tasks, randomly mixed across categories. + + ``source`` overrides the HF stream with an iterable of raw row dicts (tests). + We over-stream a buffer and shuffle it so the yielded tasks are a random mix + of the source's categories rather than a single leading shard. + """ + if source is None: + from datasets import load_dataset # lazy: optional dep / network + + source = load_dataset(GENERALTHOUGHT_ID, split="train", streaming=True) + + rng = random.Random(seed) + # Buffer a generous multiple of n so the shuffle actually mixes categories, + # but stay bounded for the multi-hundred-K streams. + buf_cap = max(n * 6, 6000) + buf: List[Task] = [] + for i, row in enumerate(source): + task = _normalize_generalthought(dict(row), index=i) + if task is None: + continue + buf.append(task) + if len(buf) >= buf_cap: + break + rng.shuffle(buf) + for task in buf[:n]: + yield task + + +# --- OpenThoughts3 ------------------------------------------------------------ + +_BOXED_RE = re.compile(r"\\boxed\{") + + +def _extract_boxed(text: str) -> Optional[str]: + """Return the contents of the last ``\\boxed{...}`` in ``text`` (brace-balanced).""" + last = None + for m in _BOXED_RE.finditer(text): + start = m.end() + depth = 1 + i = start + while i < len(text) and depth > 0: + c = text[i] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + i += 1 + if depth == 0: + last = text[start : i - 1].strip() + return last + + +def _conversation_qa(conversations: Any) -> Optional[tuple[str, str]]: + """Pull (question, gold_answer) from an OpenThoughts ``conversations`` list.""" + if not isinstance(conversations, list): + return None + question = "" + full = "" + for turn in conversations: + if not isinstance(turn, dict): + continue + who = str(turn.get("from") or "").lower() + val = str(turn.get("value") or "") + if who in ("human", "user") and not question: + question = val.strip() + elif who in ("gpt", "assistant"): + full = val + if not question or not full: + return None + # Drop the … trace; keep the post-reasoning solution. + visible = re.sub(r"(?s).*?", "", full).strip() + body = visible or full + answer = _extract_boxed(body) or _extract_boxed(full) or body.strip() + if not answer: + return None + return question, answer + + +def _normalize_openthoughts(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]: + qa = _conversation_qa(row.get("conversations")) + if qa is None: + return None + question, answer = qa + if not question or not answer: + return None + domain = str(row.get("domain") or "unknown").strip().lower() or "unknown" + task_id = str(row.get("id") or row.get("source") or f"openthoughts-{index}") + f"-{index}" + return Task(task_id=task_id, question=question, answer=answer, domain=domain) + + +def load_openthoughts( + *, + n_code: int = 2000, + n_math: int = 2000, + n_science: int = 2000, + seed: int = 42, + source: Optional[Iterable[Dict[str, Any]]] = None, +) -> Iterator[Task]: + """Yield OpenThoughts3 tasks balanced across code / math / science. + + The HF stream is sharded by domain (all-code first), so we route rows into + per-domain buffers and stop once every quota is met. ``source`` overrides the + stream with raw row dicts (tests). + """ + if source is None: + from datasets import load_dataset # lazy: optional dep / network + + source = load_dataset(OPENTHOUGHTS_ID, split="train", streaming=True) + + quotas = {"code": n_code, "math": n_math, "science": n_science} + bufs: Dict[str, List[Task]] = {k: [] for k in quotas} + for i, row in enumerate(source): + task = _normalize_openthoughts(dict(row), index=i) + if task is None: + continue + dom = task.domain + if dom in bufs and len(bufs[dom]) < quotas[dom]: + bufs[dom].append(task) + if all(len(bufs[k]) >= quotas[k] for k in quotas): + break + + rng = random.Random(seed) + out: List[Task] = [] + for k in quotas: + out.extend(bufs[k]) + rng.shuffle(out) + for task in out: + yield task + + +# --- combined SFT / GRPO sets ------------------------------------------------- + + +def load_sft_tasks(*, seed: int = 42) -> List[Task]: + """The 8K cold-start SFT set: 2K GeneralThought + 2K code + 2K math + 2K science.""" + tasks: List[Task] = [] + tasks.extend(load_generalthought(n=2000, seed=seed)) + tasks.extend(load_openthoughts(n_code=2000, n_math=2000, n_science=2000, seed=seed)) + rng = random.Random(seed) + rng.shuffle(tasks) + return tasks + + +def load_grpo_prompts(*, n: int = 30000, seed: int = 42) -> List[Task]: + """Pool ``n`` unique prompts from the same datasets, deduped by question. + + We draw a roughly even split from GeneralThought and OpenThoughts3 (code / + math / science), dedupe on the question text, and cap at ``n``. + """ + per = max(n // 4 + 1, 1) + pool: List[Task] = [] + pool.extend(load_generalthought(n=per, seed=seed)) + pool.extend( + load_openthoughts(n_code=per, n_math=per, n_science=per, seed=seed) + ) + + rng = random.Random(seed) + rng.shuffle(pool) + + seen: set[str] = set() + out: List[Task] = [] + for task in pool: + key = task.question.strip() + if not key or key in seen: + continue + seen.add(key) + out.append(task) + if len(out) >= n: + break + return out + + +__all__ = [ + "GENERALTHOUGHT_ID", + "OPENTHOUGHTS_ID", + "Task", + "load_generalthought", + "load_grpo_prompts", + "load_openthoughts", + "load_sft_tasks", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/hotpotqa.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/hotpotqa.py new file mode 100644 index 00000000..36b2f92b --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/hotpotqa.py @@ -0,0 +1,155 @@ +"""Loader + answer verifier for HotpotQA (``hotpotqa/hotpot_qa``). + +HotpotQA is multi-hop Wikipedia QA: each task is a question with a short +factoid ``answer`` and a difficulty ``level`` (easy/medium/hard). We use it as +the orchestrator SFT cold-start source because (a) answers are exact-match +verifiable, so rejection sampling needs **no LLM judge**, and (b) solving a +question requires ``web_search`` + multi-step reasoning, which is exactly the +local↔cloud routing signal we want to teach. + +We load the ``fullwiki`` config (the agent must retrieve from all of Wikipedia +rather than being handed the gold paragraphs); we only keep ``question`` / +``answer`` / ``level`` / ``type`` since the orchestrator does its own retrieval. + +``load_hotpotqa`` streams via the HuggingFace ``datasets`` library; tests pass +``source=`` an iterable of raw row dicts so normalization is exercised offline. +""" + +from __future__ import annotations + +import re +import string +from dataclasses import dataclass +from typing import Any, Dict, Iterable, Iterator, List, Optional + +DATASET_ID = "hotpotqa/hotpot_qa" +CONFIG = "fullwiki" +SPLIT = "validation" # train/validation carry gold answers; test does not + + +@dataclass +class HotpotTask: + task_id: str + question: str + answer: str + level: str = "" # easy | medium | hard + qtype: str = "" # comparison | bridge + + # Parity with ToolScaleTask so the rejection-sampling loop is dataset-agnostic. + @property + def instruction(self) -> str: + return self.question + + @property + def domain(self) -> str: + return f"hotpotqa/{self.level or 'unknown'}" + + +def normalize_row(row: Dict[str, Any], *, index: int = 0) -> Optional[HotpotTask]: + """Turn one raw HotpotQA row into a :class:`HotpotTask` (pure).""" + question = str(row.get("question") or "").strip() + answer = str(row.get("answer") or "").strip() + if not question or not answer: + return None + task_id = str(row.get("id") or row.get("_id") or f"hotpotqa-{index}") + return HotpotTask( + task_id=task_id, + question=question, + answer=answer, + level=str(row.get("level") or ""), + qtype=str(row.get("type") or ""), + ) + + +def load_hotpotqa( + *, + max_tasks: Optional[int] = None, + split: str = SPLIT, + config: str = CONFIG, + source: Optional[Iterable[Dict[str, Any]]] = None, +) -> Iterator[HotpotTask]: + """Yield normalized HotpotQA tasks. + + ``source`` overrides the HF stream with an iterable of raw row dicts (tests). + When ``source`` is None, streams ``hotpotqa/hotpot_qa`` via ``datasets``. + """ + if source is None: + from datasets import load_dataset # lazy: optional dep / network + + source = load_dataset(DATASET_ID, config, split=split, streaming=True) + + n = 0 + for i, row in enumerate(source): + if max_tasks is not None and n >= max_tasks: + break + task = normalize_row(dict(row), index=i) + if task is None: + continue + yield task + n += 1 + + +# --- answer verification (official HotpotQA-style normalization) ------------- + + +def _normalize_answer(s: str) -> str: + """Lowercase, strip punctuation / articles / extra whitespace (SQuAD/HotpotQA).""" + s = s.lower() + s = "".join(ch for ch in s if ch not in set(string.punctuation)) + s = re.sub(r"\b(a|an|the)\b", " ", s) + return " ".join(s.split()) + + +def _f1(pred: str, gold: str) -> float: + pt, gt = _normalize_answer(pred).split(), _normalize_answer(gold).split() + if not pt or not gt: + return float(pt == gt) + common: Dict[str, int] = {} + for w in pt: + if w in gt: + common[w] = common.get(w, 0) + 1 + num_same = sum(common.values()) + if num_same == 0: + return 0.0 + precision = num_same / len(pt) + recall = num_same / len(gt) + return 2 * precision * recall / (precision + recall) + + +def answer_matches(prediction: str, gold: str, *, f1_threshold: float = 0.6) -> bool: + """True if ``prediction`` (free-form) contains / matches the gold answer. + + The orchestrator answers in prose, so we accept (a) the normalized gold as a + whole-word substring of the prediction, or (b) token-F1 >= threshold. This is + the dependency-free, no-LLM verifier that makes rejection sampling automatic. + """ + np, ng = _normalize_answer(prediction), _normalize_answer(gold) + if not ng: + return False + if f" {ng} " in f" {np} ": + return True + return _f1(prediction, gold) >= f1_threshold + + +def make_verifier(f1_threshold: float = 0.6): + """Return a ``verify_fn(task, rollout) -> bool`` for the rejection sampler.""" + + def verify(task: HotpotTask, rollout: Any) -> bool: + ans = getattr(rollout, "final_answer", "") or "" + if not ans.strip(): + return False + return answer_matches(ans, task.answer, f1_threshold=f1_threshold) + + return verify + + +__all__ = [ + "CONFIG", + "DATASET_ID", + "HotpotTask", + "SPLIT", + "answer_matches", + "load_hotpotqa", + "make_verifier", + "normalize_row", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py deleted file mode 100644 index 86b3089b..00000000 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Render a canonical ADP ``Episode`` under each local↔cloud paradigm. - -Each renderer assigns a model **tier** to every step according to that -paradigm's policy, attaches the estimated telemetry (cost/energy/latency/power) -to the step's observation, and predicts whether the rendering would still solve -the task. Selection (``select.py``) then keeps the cheapest predicted-correct -rendering per task. - -These are the cold-start *tiering* scaffolds of the real -``agents/hybrid/*`` paradigms — same names, same local-first spirit, no live -execution. Swapping in real runs is the v2 item in the design doc. -""" - -from __future__ import annotations - -import copy -from dataclasses import dataclass, field -from typing import Callable, List - -from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import ( - Difficulty, - Tier, - covers, - min_covering_tier, - step_difficulty, - tier_name, - tier_telemetry, -) -from openjarvis.learning.intelligence.orchestrator.types import Episode - -# Detected in an observation -> the demonstration had to retry -> harder task. -_RETRY_MARKERS = ( - "your answer is wrong", - "incorrect", - "traceback", - "error:", - "failed", - "try again", -) - - -@dataclass -class RenderedEpisode: - """An ``Episode`` re-tiered under one paradigm, plus its prediction.""" - - paradigm: str - episode: Episode - """Deep copy with per-step tier in ``action.tool_name`` and telemetry filled.""" - - step_tiers: List[str] = field(default_factory=list) - """Tier key (``local``/``mid``/``frontier``/``search``) chosen per step.""" - - predicted_correct: bool = False - - -def _difficulties(episode: Episode) -> List[Difficulty]: - """Per-step difficulty, with a trajectory-wide retry bump.""" - retry = any( - any(m in (step.observation.content or "").lower() for m in _RETRY_MARKERS) - for step in episode.steps - ) - out: List[Difficulty] = [] - for step in episode.steps: - kind = step.action.tool_name - out.append(step_difficulty(kind, step.action.tool_input, retry_signal=retry)) - return out - - -def _apply( - episode: Episode, - paradigm: str, - tier_keys: List[str], - predicted_correct: bool, -) -> RenderedEpisode: - """Build a RenderedEpisode: stamp each step with its tier + telemetry.""" - ep = copy.deepcopy(episode) - ep.total_cost_usd = 0.0 - ep.total_energy_joules = 0.0 - ep.total_latency_seconds = 0.0 - ep.total_tokens = 0 - ep.max_power_watts = 0.0 - - for step, tier_key in zip(ep.steps, tier_keys): - tel = tier_telemetry(tier_key) - step.observation.cost_usd = tel["cost_usd"] - step.observation.energy_joules = tel["energy_joules"] - step.observation.latency_seconds = tel["latency_seconds"] - step.observation.power_watts = tel["power_watts"] - step.observation.tokens = int(tel["tokens"]) - # Record the routing decision on the action for the serializer. - step.action.thought = _thought_for(tier_key, step.action.tool_name) - ep.total_cost_usd += tel["cost_usd"] - ep.total_energy_joules += tel["energy_joules"] - ep.total_latency_seconds += tel["latency_seconds"] - ep.total_tokens += int(tel["tokens"]) - ep.max_power_watts = max(ep.max_power_watts, tel["power_watts"]) - - ep.correct = predicted_correct - ep.metadata["paradigm"] = paradigm - return RenderedEpisode( - paradigm=paradigm, - episode=ep, - step_tiers=tier_keys, - predicted_correct=predicted_correct, - ) - - -def _thought_for(tier_key: str, kind: str) -> str: - if kind == "search" or tier_key == "search": - return "This step needs retrieval; route to web_search." - reason = { - "local": "Cheap and on-device; the local model can handle this step.", - "mid": "Beyond easy; escalate to a cheap mid-tier cloud model.", - "frontier": "Hard step; escalate to the frontier cloud model.", - }[tier_key] - return reason - - -def _tier_key_for_step(kind: str, tier: Tier) -> str: - return "search" if kind == "search" else tier_name(tier) - - -# --- paradigm renderers ----------------------------------------------------- - - -def render_baseline_local(episode: Episode) -> RenderedEpisode: - diffs = _difficulties(episode) - tiers = [ - _tier_key_for_step(s.action.tool_name, Tier.LOCAL) for s in episode.steps - ] - predicted = all(d == Difficulty.EASY for d in diffs) - return _apply(episode, "baseline_local", tiers, predicted) - - -def render_baseline_cloud(episode: Episode) -> RenderedEpisode: - tiers = [ - _tier_key_for_step(s.action.tool_name, Tier.FRONTIER) for s in episode.steps - ] - return _apply(episode, "baseline_cloud", tiers, True) - - -def render_advisor(episode: Episode) -> RenderedEpisode: - """Local executor, frontier rewrite of the final answer.""" - diffs = _difficulties(episode) - tiers: List[str] = [] - n = len(episode.steps) - for i, s in enumerate(episode.steps): - tier = Tier.FRONTIER if i == n - 1 else Tier.LOCAL - tiers.append(_tier_key_for_step(s.action.tool_name, tier)) - predicted = all(d == Difficulty.EASY for d in diffs[:-1]) if n > 1 else True - return _apply(episode, "advisor", tiers, predicted) - - -def render_toolorchestra(episode: Episode) -> RenderedEpisode: - """Local-first per-step: assign each step its minimum covering tier.""" - diffs = _difficulties(episode) - tiers = [ - _tier_key_for_step(s.action.tool_name, min_covering_tier(d)) - for s, d in zip(episode.steps, diffs) - ] - predicted = all( - covers(min_covering_tier(d), d) for d in diffs - ) # True by construction - return _apply(episode, "toolorchestra", tiers, predicted) - - -PARADIGMS: dict[str, Callable[[Episode], RenderedEpisode]] = { - "baseline_local": render_baseline_local, - "baseline_cloud": render_baseline_cloud, - "advisor": render_advisor, - "toolorchestra": render_toolorchestra, -} - - -def render_all(episode: Episode) -> List[RenderedEpisode]: - """Render ``episode`` under every paradigm.""" - return [render(episode) for render in PARADIGMS.values()] - - -__all__ = [ - "PARADIGMS", - "RenderedEpisode", - "render_advisor", - "render_all", - "render_baseline_cloud", - "render_baseline_local", - "render_toolorchestra", -] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py index 605efecf..62d30130 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py @@ -23,22 +23,27 @@ from collections import Counter from pathlib import Path from typing import Callable, Iterable, List, Optional +from typing import Any + from openjarvis.agents.hybrid.expert_registry import ExpertTool from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout -from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import ( - ToolScaleTask, -) from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import ( trajectory_to_record, ) logger = logging.getLogger(__name__) -RolloutFn = Callable[[ToolScaleTask], Optional[UnifiedRollout]] -VerifyFn = Callable[[ToolScaleTask, UnifiedRollout], bool] +# The sampler only touches ``task.task_id`` / ``task.instruction`` / ``task.domain`` +# (via ``trajectory_to_record``), so any task dataclass works — ToolScaleTask or the +# reasoning ``datasets.Task``. ``gold_coverage_verify`` additionally needs +# ``gold_action_names()`` (ToolScale-only), but callers using ``datasets.Task`` pass +# their own ``verify_fn`` (e.g. ``verify.make_verifier()``) instead. +TaskLike = Any +RolloutFn = Callable[[TaskLike], Optional[UnifiedRollout]] +VerifyFn = Callable[[TaskLike, UnifiedRollout], bool] -def gold_coverage_verify(task: ToolScaleTask, rollout: UnifiedRollout) -> bool: +def gold_coverage_verify(task: TaskLike, rollout: UnifiedRollout) -> bool: """Dependency-free proxy verifier: trajectory must (a) produce a non-empty answer and (b) call tools covering every golden action name. @@ -57,7 +62,7 @@ def gold_coverage_verify(task: ToolScaleTask, rollout: UnifiedRollout) -> bool: def generate_sft_dataset( out_path: str, *, - tasks: Iterable[ToolScaleTask], + tasks: Iterable[TaskLike], tools: List[ExpertTool], rollout_fn: RolloutFn, verify_fn: VerifyFn = gold_coverage_verify, diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py deleted file mode 100644 index 07097474..00000000 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Pick the best paradigm rendering per task. - -The label for SFT is the *cheapest strategy that still solves the task*: among -the predicted-correct renderings, keep the one with the highest multi-objective -reward (accuracy minus cost/energy/latency/power). This is the local-first -thesis expressed as a training target. -""" - -from __future__ import annotations - -from typing import List, Optional - -from openjarvis.learning.intelligence.orchestrator.reward import ( - MultiObjectiveReward, - Normalizers, - RewardWeights, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import ( - RenderedEpisode, -) - -_DEFAULT_REWARD = MultiObjectiveReward(RewardWeights(), Normalizers()) - - -def select_best( - renderings: List[RenderedEpisode], - *, - reward: Optional[MultiObjectiveReward] = None, -) -> Optional[RenderedEpisode]: - """Return the highest-reward predicted-correct rendering, or ``None``. - - ``None`` means no paradigm was predicted to solve the task — the task is - dropped from the SFT set rather than teaching a wrong trajectory. - """ - scorer = reward or _DEFAULT_REWARD - correct = [r for r in renderings if r.predicted_correct] - if not correct: - return None - return max(correct, key=lambda r: scorer.compute(r.episode)) - - -__all__ = ["select_best"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py deleted file mode 100644 index f9bbd1ad..00000000 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Serialize a winning :class:`RenderedEpisode` into a ``conversations`` record. - -Output schema matches what -:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset` -already consumes:: - - {"conversations": [{"role": "system"|"user"|"assistant"|"tool", "content": ...}], - "paradigm": ..., "reward": ..., "metrics": {...}} - -The assistant turns use the canonical THOUGHT/TOOL/INPUT format from -``prompt_registry``; the routing tool encodes the tier the orchestrator chose -(``local_model`` / ``mid_model`` / ``frontier_model`` / ``web_search``). -""" - -from __future__ import annotations - -from typing import Any, Dict - -from openjarvis.learning.intelligence.orchestrator.prompt_registry import ( - build_system_prompt, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import ( - RenderedEpisode, -) - -# Routing vocabulary the orchestrator learns to emit. -TIER_TOOL = { - "local": "local_model", - "mid": "mid_model", - "frontier": "frontier_model", - "search": "web_search", -} -ROUTING_TOOLS = ["local_model", "mid_model", "frontier_model", "web_search"] - -_SYSTEM_PROMPT = build_system_prompt(ROUTING_TOOLS) - - -def to_record(rendered: RenderedEpisode, *, reward: float = 0.0) -> Dict[str, Any]: - """Convert a winning rendering into one SFT JSONL record.""" - ep = rendered.episode - conversations: list[dict[str, str]] = [ - {"role": "system", "content": _SYSTEM_PROMPT}, - {"role": "user", "content": ep.initial_prompt}, - ] - - n = len(ep.steps) - for i, (step, tier_key) in enumerate(zip(ep.steps, rendered.step_tiers)): - tool = TIER_TOOL.get(tier_key, "local_model") - if step.action.is_final_answer or i == n - 1: - conversations.append( - { - "role": "assistant", - "content": ( - f"THOUGHT: {step.action.thought}\n" - f"TOOL: {tool}\n" - f"INPUT: {step.action.tool_input}" - ), - } - ) - conversations.append( - {"role": "tool", "name": tool, "content": step.observation.content} - ) - conversations.append( - { - "role": "assistant", - "content": ( - "THOUGHT: The result answers the task.\n" - f"FINAL_ANSWER: {ep.final_answer}" - ), - } - ) - else: - conversations.append( - { - "role": "assistant", - "content": ( - f"THOUGHT: {step.action.thought}\n" - f"TOOL: {tool}\n" - f"INPUT: {step.action.tool_input}" - ), - } - ) - conversations.append( - {"role": "tool", "name": tool, "content": step.observation.content} - ) - - return { - "conversations": conversations, - "task_id": ep.task_id, - "paradigm": rendered.paradigm, - "reward": reward, - "metrics": { - "cost_usd": ep.total_cost_usd, - "energy_joules": ep.total_energy_joules, - "latency_seconds": ep.total_latency_seconds, - "tokens": ep.total_tokens, - "max_power_watts": ep.max_power_watts, - "num_steps": n, - }, - } - - -__all__ = ["ROUTING_TOOLS", "TIER_TOOL", "to_record"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py deleted file mode 100644 index 4632a262..00000000 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Model tiers, per-tier telemetry estimates, and step-difficulty heuristics. - -This is the *one estimate* of the cold-start (see the design doc): because we -do not execute models, we approximate "what would this step cost on tier X?" -with flat per-tier telemetry, and "how hard is this step?" with a heuristic over -the ADP step kind and the trajectory's retry signals. - -Costs reuse the authoritative hybrid pricing table so the orchestrator is -trained against the same dollars the paradigm harness charges. -""" - -from __future__ import annotations - -from enum import IntEnum - -from openjarvis.agents.hybrid import _prices - -# Representative model per tier (drives the $ side of telemetry). -TIER_MODEL: dict[str, str] = { - "local": "qwen3:8b", # local vLLM -> unknown to PRICES -> $0 - "mid": "gemini-2.5-flash", - "frontier": "claude-sonnet-4-6", - "search": "qwen3:8b", # retrieval is a tool call, not a model tier -} - -# Rough token budgets per step kind, for the cost estimate. -_TOKENS_IN = 800 -_TOKENS_OUT = {"local": 400, "mid": 400, "frontier": 600, "search": 80} - -# Non-$ telemetry per step (joules / seconds / watts). Local work burns some -# on-device energy/power but no dollars; cloud work is ~free locally but costs -# dollars and adds latency. Tuned (against reward.py's default weights and -# normalizers) so the local-first ordering holds: on a step a tier can cover, -# local beats cloud, and per-step escalation beats all-frontier. These are the -# tunable knobs of the cold-start estimate — see the design doc. -_ENERGY_J = {"local": 5.0, "mid": 0.5, "frontier": 0.5, "search": 0.2} -_LATENCY_S = {"local": 1.5, "mid": 1.5, "frontier": 4.0, "search": 1.0} -_POWER_W = {"local": 20.0, "mid": 5.0, "frontier": 5.0, "search": 5.0} -_SEARCH_COST_USD = 0.001 - - -class Tier(IntEnum): - """Model capability tier; ordered so a higher tier covers a harder step.""" - - LOCAL = 0 - MID = 1 - FRONTIER = 2 - - -class Difficulty(IntEnum): - """Estimated step difficulty; compared against :class:`Tier` rank.""" - - EASY = 0 - MED = 1 - HARD = 2 - - -_TIER_NAME = {Tier.LOCAL: "local", Tier.MID: "mid", Tier.FRONTIER: "frontier"} - - -def tier_name(tier: Tier) -> str: - return _TIER_NAME[tier] - - -def tier_telemetry(tier_key: str) -> dict[str, float]: - """Estimated ``(cost, energy, latency, power, tokens)`` for one step at a tier. - - ``tier_key`` is one of ``local``/``mid``/``frontier``/``search``. - """ - out_tokens = _TOKENS_OUT[tier_key] - if tier_key == "search": - cost = _SEARCH_COST_USD - else: - cost = _prices.cost(TIER_MODEL[tier_key], _TOKENS_IN, out_tokens) - return { - "cost_usd": cost, - "energy_joules": _ENERGY_J[tier_key], - "latency_seconds": _LATENCY_S[tier_key], - "power_watts": _POWER_W[tier_key], - "tokens": float(_TOKENS_IN + out_tokens), - } - - -def step_difficulty( - kind: str, content: str, *, retry_signal: bool = False -) -> Difficulty: - """Heuristic difficulty for a canonical step. - - - ``code`` actions are at least MED (real execution / logic). - - very long actions, or any step in a trajectory that hit a retry signal - ("your answer is wrong" etc.), bump to HARD. - - everything else (plain messages / short reasoning) is EASY. - """ - base = Difficulty.MED if kind == "code" else Difficulty.EASY - if retry_signal or len(content) > 1500: - return Difficulty.HARD - return base - - -def covers(tier: Tier, difficulty: Difficulty) -> bool: - """True iff ``tier`` is capable enough for a step of this ``difficulty``.""" - return int(tier) >= int(difficulty) - - -def min_covering_tier(difficulty: Difficulty) -> Tier: - """The cheapest tier that still covers ``difficulty`` (local-first optimum).""" - return Tier(int(difficulty)) - - -__all__ = [ - "Difficulty", - "Tier", - "TIER_MODEL", - "covers", - "min_covering_tier", - "step_difficulty", - "tier_name", - "tier_telemetry", -] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py index 6f9761c4..cdcb61f8 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py @@ -17,19 +17,15 @@ from typing import Any, Dict, List from openjarvis.agents.hybrid.expert_registry import ExpertTool, build_tool_specs from openjarvis.agents.hybrid.toolorchestra.rollout import ( - RL_ORCHESTRATOR_SYS, UnifiedRollout, + build_system_prompt, tool_call_tag, ) def _system_prompt(tools: List[ExpertTool]) -> str: - specs = build_tool_specs(tools) - return ( - RL_ORCHESTRATOR_SYS - + "\n\nAvailable tools (call one per turn, or answer directly):\n" - + json.dumps(specs, indent=2) - ) + # Faithful ToolOrchestra system prompt (paper's prepare_sft_data.py format). + return build_system_prompt(build_tool_specs(tools)) def trajectory_to_record( @@ -44,7 +40,7 @@ def trajectory_to_record( """Convert a passing :class:`UnifiedRollout` into one SFT JSONL record.""" conversations: List[Dict[str, str]] = [ {"role": "system", "content": _system_prompt(tools)}, - {"role": "user", "content": f"Problem: {question}\n\nChoose an appropriate tool."}, + {"role": "user", "content": f"Problem: {question}"}, ] for turn in rollout.turns: diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py new file mode 100644 index 00000000..31c88c87 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py @@ -0,0 +1,254 @@ +"""Correctness verifier for orchestrator reasoning tasks (:mod:`.datasets`). + +The verifier dispatches on :attr:`Task.domain`: + +- **math** — normalize and compare. Extract ``\\boxed{...}`` from both sides, + try ``sympy`` for symbolic / numeric equality, and fall back to a normalized + string / number match. No network. +- **code** — best-effort: if a short expected output / answer exists, do a + normalized substring match; otherwise defer to the LLM judge. (We don't run + arbitrary code here.) +- **science / medical / chat / misc / unknown** — Gemini LLM judge given the + question + gold answer + candidate, asked for ``PASS`` / ``FAIL``. When no + ``GEMINI_API_KEY`` is set (e.g. offline tests) it falls back to a normalized + string / token-F1 >= 0.6 match. + +The OpenAI key is dead, so the LLM judge talks to Gemini through its +OpenAI-compatible endpoint. Normalization helpers are copied (not imported) +from ``hotpotqa.py`` to keep this module self-contained. +""" + +from __future__ import annotations + +import os +import re +import string +from typing import Any, Callable, Dict, Optional + +from .datasets import Task + +GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/" +GEMINI_MODEL = "gemini-2.5-flash" + + +# --- normalization (copied from hotpotqa.py — keep self-contained) ----------- + + +def _normalize_answer(s: str) -> str: + """Lowercase, strip punctuation / articles / extra whitespace (SQuAD/HotpotQA).""" + s = s.lower() + s = "".join(ch for ch in s if ch not in set(string.punctuation)) + s = re.sub(r"\b(a|an|the)\b", " ", s) + return " ".join(s.split()) + + +def _f1(pred: str, gold: str) -> float: + pt, gt = _normalize_answer(pred).split(), _normalize_answer(gold).split() + if not pt or not gt: + return float(pt == gt) + common: Dict[str, int] = {} + for w in pt: + if w in gt: + common[w] = common.get(w, 0) + 1 + num_same = sum(common.values()) + if num_same == 0: + return 0.0 + precision = num_same / len(pt) + recall = num_same / len(gt) + return 2 * precision * recall / (precision + recall) + + +def _string_or_f1(prediction: str, gold: str, *, f1_threshold: float = 0.6) -> bool: + """Normalized whole-word substring OR token-F1 >= threshold.""" + np_, ng = _normalize_answer(prediction), _normalize_answer(gold) + if not ng: + return False + if f" {ng} " in f" {np_} ": + return True + return _f1(prediction, gold) >= f1_threshold + + +# --- math -------------------------------------------------------------------- + +_BOXED_RE = re.compile(r"\\boxed\{") +_NUM_RE = re.compile(r"-?\d+(?:\.\d+)?") + + +def _extract_boxed(text: str) -> Optional[str]: + """Return the contents of the last ``\\boxed{...}`` (brace-balanced).""" + last = None + for m in _BOXED_RE.finditer(text): + start = m.end() + depth = 1 + i = start + while i < len(text) and depth > 0: + c = text[i] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + i += 1 + if depth == 0: + last = text[start : i - 1].strip() + return last + + +def _clean_math(s: str) -> str: + s = s.strip() + boxed = _extract_boxed(s) + if boxed is not None: + s = boxed + # Common LaTeX wrappers / delimiters. + s = s.replace("$", "").replace("\\!", "").replace("\\,", "").replace("\\;", "") + s = s.replace("\\left", "").replace("\\right", "") + s = s.strip().strip("$ ") + return s + + +def _to_float(s: str) -> Optional[float]: + s = s.replace(",", "").strip() + try: + return float(s) + except (ValueError, TypeError): + pass + m = _NUM_RE.search(s) + if m: + try: + return float(m.group(0)) + except ValueError: + return None + return None + + +def _math_equal(prediction: str, gold: str) -> bool: + pred = _clean_math(prediction) + g = _clean_math(gold) + if not g: + return False + + # Exact normalized-string shortcut. + if pred.replace(" ", "") == g.replace(" ", "") and pred: + return True + + # Numeric comparison (handles "42" vs "42.0" vs trailing prose). + pf, gf = _to_float(pred), _to_float(g) + if pf is not None and gf is not None: + if abs(pf - gf) <= 1e-6 * max(1.0, abs(gf)): + return True + + # Symbolic / exact equality via sympy when available. + try: + import sympy + from sympy.parsing.latex import parse_latex # noqa: F401 (optional) + + def _parse(x: str): + try: + return sympy.sympify(x.replace("^", "**")) + except Exception: + try: + return sympy.parsing.latex.parse_latex(x) + except Exception: + return None + + pe, ge = _parse(pred), _parse(g) + if pe is not None and ge is not None: + try: + if sympy.simplify(pe - ge) == 0: + return True + except Exception: + pass + except Exception: + pass + + # Last resort: whole-word substring of the gold in the prediction. + return _string_or_f1(prediction, g, f1_threshold=0.9) + + +# --- LLM judge (Gemini, OpenAI-compatible) ----------------------------------- + + +def _gemini_judge(task: Task, prediction: str) -> Optional[bool]: + """Ask Gemini PASS/FAIL. Returns None if unavailable (caller falls back).""" + api_key = os.environ.get("GEMINI_API_KEY") + if not api_key: + return None + try: + from openai import OpenAI + except Exception: + return None + + try: + client = OpenAI(api_key=api_key, base_url=GEMINI_BASE_URL) + prompt = ( + "You are grading a candidate answer against a gold reference.\n" + "Reply with exactly one word: PASS if the candidate is correct and " + "consistent with the gold answer, otherwise FAIL.\n\n" + f"Question:\n{task.question}\n\n" + f"Gold answer:\n{task.answer}\n\n" + f"Candidate answer:\n{prediction}\n\n" + "Verdict (PASS or FAIL):" + ) + resp = client.chat.completions.create( + model=GEMINI_MODEL, + messages=[{"role": "user", "content": prompt}], + temperature=0.0, + max_tokens=8, + ) + verdict = (resp.choices[0].message.content or "").strip().upper() + if "PASS" in verdict: + return True + if "FAIL" in verdict: + return False + return None + except Exception: + return None + + +# --- public API -------------------------------------------------------------- + + +def verify_answer(task: Task, prediction: str) -> bool: + """Domain-dispatched correctness check for ``prediction`` against ``task.answer``.""" + pred = (prediction or "").strip() + if not pred or not (task.answer or "").strip(): + return False + + domain = (task.domain or "").lower() + + if domain == "math": + return _math_equal(pred, task.answer) + + if domain == "code": + # Best-effort: short gold -> normalized substring; otherwise LLM judge. + gold = task.answer.strip() + if len(gold) <= 200: + if _string_or_f1(pred, gold, f1_threshold=0.8): + return True + judged = _gemini_judge(task, pred) + if judged is not None: + return judged + return _string_or_f1(pred, gold, f1_threshold=0.6) + + # science / medical / chat / misc / unknown -> LLM judge, then F1 fallback. + judged = _gemini_judge(task, pred) + if judged is not None: + return judged + return _string_or_f1(pred, task.answer, f1_threshold=0.6) + + +def make_verifier() -> Callable[[Any, Any], bool]: + """Return ``verify_fn(task, rollout) -> bool`` for the rejection sampler.""" + + def verify(task: Any, rollout: Any) -> bool: + ans = getattr(rollout, "final_answer", "") or "" + return verify_answer(task, ans) + + return verify + + +__all__ = [ + "GEMINI_BASE_URL", + "GEMINI_MODEL", + "make_verifier", + "verify_answer", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py index 9cd39522..9e467588 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py @@ -52,7 +52,7 @@ class OrchestratorSFTConfig: """Configuration for orchestrator SFT training.""" # Model - model_name: str = "Qwen/Qwen3-1.7B" + model_name: str = "Qwen/Qwen3.5-9B" max_seq_length: int = 4096 # Training @@ -63,21 +63,27 @@ class OrchestratorSFTConfig: warmup_ratio: float = 0.1 max_grad_norm: float = 1.0 - # Trace generation - teacher_engine_key: str = "" - teacher_model: str = "" - traces_per_query: int = 2 - max_attempts_per_trace: int = 3 + # Trace generation by base Qwen3-8B self-sampling (v1 cold-start). The + # orchestrator IS the local Qwen3-8B served over an OpenAI-compatible vLLM + # endpoint; we roll it out over load_sft_tasks() and keep correct trajectories. + # For v2, point orchestrator_endpoint/model at the v1 checkpoint's server. + orchestrator_endpoint: str = "http://localhost:8001/v1" + orchestrator_model: str = "qwen3-8b" + orchestrator_api_key: str = "EMPTY" generation_temperature: float = 0.7 # Data source - trace_cache_path: str = "data/orchestrator_sft_traces.jsonl" + trace_cache_path: str = "data/orchestrator_sft_v1.jsonl" regenerate_traces: bool = False - # ADP cold-start trace synthesis (see sft_data/build.py). Empty configs - # falls back to sft_data.adp_loader.DEFAULT_CONFIGS. - adp_configs: str = "" - distill_max_tasks: int = 2000 + # Rejection-sampling cold-start over load_sft_tasks (see sft_data/reject_sample.py). + distill_max_tasks: int = 0 # 0 / None -> all of load_sft_tasks() + samples_per_task: int = 8 + max_keep_per_task: int = 1 + max_rollout_turns: int = 8 + # Optional local OSS model endpoints, mapping full model id (e.g. + # "Qwen/Qwen3.5-9B") -> vLLM base_url; omitted when unset. + local_endpoints: Dict[str, str] = field(default_factory=dict) # Checkpoint checkpoint_dir: str = "checkpoints/orchestrator_sft" @@ -272,31 +278,74 @@ class OrchestratorSFTTrainer: ) def _generate_traces(self) -> None: - """Generate SFT traces from the ADP corpus (no GPU / no API keys). + """Generate SFT traces by base Qwen3-8B self-sampling (v1 cold-start). - Builds the THOUGHT/TOOL/INPUT ``conversations`` JSONL by re-tiering - NeuLab ADP trajectories through the local↔cloud paradigms and keeping - the cheapest predicted-correct rendering per task. See - ``sft_data/build.py`` and the cold-start design doc. + The orchestrator IS the local Qwen3-8B served over an OpenAI-compatible + vLLM endpoint (``orchestrator_endpoint``/``orchestrator_model``). We roll + it out over ``load_sft_tasks()`` on the fixed orchestrator tool catalog, + verify each trajectory's final answer against the gold + (``verify.make_verifier``), keep the cheapest-correct trajectory per task, + and serialize them into the ```` ``conversations`` JSONL the SFT + dataset consumes. For v2, point the endpoint/model at the v1 checkpoint's + server. See ``sft_data/reject_sample.py`` and the + ``scripts/orchestrator/build_orchestrator_sft.py`` driver. """ - from openjarvis.learning.intelligence.orchestrator.sft_data.build import ( - build_sft_dataset, + from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog + from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout + from openjarvis.agents.hybrid.toolorchestra.unified import ( + make_call_orchestrator, + make_dispatch, + ) + from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import ( + load_sft_tasks, + ) + from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import ( + generate_sft_dataset, + ) + from openjarvis.learning.intelligence.orchestrator.sft_data.verify import ( + make_verifier, ) trace_path = Path(self.config.trace_cache_path) - configs = ( - [c.strip() for c in self.config.adp_configs.split(",") if c.strip()] - or None + tools = orchestrator_catalog( + local_endpoints=self.config.local_endpoints or None, ) + call_orch = make_call_orchestrator( + self.config.orchestrator_model, + base_url=self.config.orchestrator_endpoint, + api_key=self.config.orchestrator_api_key, + temperature=self.config.generation_temperature, + ) + dispatch = make_dispatch({}) + + def rollout_fn(task: Any) -> Any: + try: + return run_unified_rollout( + task.instruction, tools, + call_orchestrator=call_orch, dispatch=dispatch, + max_turns=self.config.max_rollout_turns, + ) + except Exception as exc: # network/key failures shouldn't kill the run + logger.warning("rollout failed for %s: %s", task.task_id, exc) + return None + try: - stats = build_sft_dataset( + tasks = load_sft_tasks() + if self.config.distill_max_tasks: + tasks = tasks[: self.config.distill_max_tasks] + stats = generate_sft_dataset( str(trace_path), - max_tasks=self.config.distill_max_tasks, - **({"configs": configs} if configs else {}), + tasks=tasks, + tools=tools, + rollout_fn=rollout_fn, + verify_fn=make_verifier(), + samples_per_task=self.config.samples_per_task, + max_keep_per_task=self.config.max_keep_per_task, + reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets top reward ) logger.info("Generated SFT traces: %s", stats) except Exception as exc: # network/datasets optional — fail soft to empty - logger.warning("ADP trace generation failed (%s); writing empty set", exc) + logger.warning("trace generation failed (%s); writing empty set", exc) trace_path.parent.mkdir(parents=True, exist_ok=True) if not trace_path.exists(): trace_path.touch() diff --git a/tests/agents/test_expert_registry.py b/tests/agents/test_expert_registry.py index c635e0ac..62222d7e 100644 --- a/tests/agents/test_expert_registry.py +++ b/tests/agents/test_expert_registry.py @@ -7,10 +7,16 @@ import random import pytest from openjarvis.agents.hybrid.expert_registry import ( + CATEGORY_BASIC, + CATEGORY_CLOUD_FRONTIER, + CATEGORY_LOCAL_OSS, ExpertTool, KIND_MODEL, + KIND_TOOL, build_tool_specs, default_catalog, + openjarvis_tool, + orchestrator_catalog, sample_tool_config, to_worker_dict, tools_by_name, @@ -22,8 +28,8 @@ def test_each_model_is_its_own_tool(): cat = default_catalog() names = {t.name for t in cat} # Distinct model tools, each with its own name. - for n in ("gpt_5", "gpt_5_mini", "qwen3_32b", "qwen2_5_coder_32b", - "llama_3_3_70b", "claude_opus"): + for n in ("gpt_5", "gpt_5_mini", "qwen3_32b", "qwen_2_5_coder_32b_instruct", + "llama_3_3_70b_instruct", "claude_opus_4_7"): assert n in names, f"missing model tool {n}" # No meta-tool / slot vocabulary leaks in. assert "answer" not in names and "enhance_reasoning" not in names @@ -37,9 +43,10 @@ def test_catalog_names_unique_and_valid(): def test_local_model_included_only_when_served(): - assert "local_model" not in {t.name for t in default_catalog()} + # Local tool is named after the served model ("qwen3:8b" -> "qwen3_8b"). + assert "qwen3_8b" not in {t.name for t in default_catalog()} cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1") - local = tools_by_name(cat)["local_model"] + local = tools_by_name(cat)["qwen3_8b"] assert local.backend_type == "vllm" assert local.base_url == "http://x/v1" assert local.price_in == 0.0 and local.price_out == 0.0 @@ -92,10 +99,121 @@ def test_price_jitter_changes_prices_reproducibly(): assert 0.5 <= f <= 1.5 +# Two cloud-frontier + four local-OSS model tools, in catalog order. +_ORCH_MODEL_NAMES = [ + "gpt_5_5", + "claude_opus_4_8", + "qwen3_5_9b", + "qwen3_6_27b_fp8", + "qwen3_5_122b_a10b_fp8", + "qwen3_5_397b_a17b_fp8", +] + +# Bridged real OpenJarvis tools (basic) appended after web_search/code_interpreter. +_ORCH_BASIC_NAMES = [ + "web_search", + "code_interpreter", + "calculator", + "shell_exec", + "file_read", + "file_write", + "http_request", + "think", + "apply_patch", + "pdf_extract", + "db_query", +] + + +def test_orchestrator_catalog_two_model_classes_plus_basics(): + cat = orchestrator_catalog() + names = [t.name for t in cat] + # 6 model tools (2 cloud_frontier + 4 local_oss) come first. + assert names[:6] == _ORCH_MODEL_NAMES + # then the basic tools. + assert set(_ORCH_BASIC_NAMES) <= set(names) + assert len(cat) == 6 + len(_ORCH_BASIC_NAMES) + by = tools_by_name(cat) + assert by["gpt_5_5"].category == CATEGORY_CLOUD_FRONTIER + assert by["claude_opus_4_8"].category == CATEGORY_CLOUD_FRONTIER + for n in ("qwen3_5_9b", "qwen3_6_27b_fp8", "qwen3_5_122b_a10b_fp8", + "qwen3_5_397b_a17b_fp8"): + assert by[n].category == CATEGORY_LOCAL_OSS + assert by[n].backend_type == "vllm" + assert by[n].price_in == 0.0 and by[n].price_out == 0.0 + + +def test_orchestrator_catalog_categories_present(): + cat = orchestrator_catalog() + cats = {t.category for t in cat} + assert cats == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS, CATEGORY_BASIC} + + +def test_orchestrator_can_drop_tools(): + cat = orchestrator_catalog(include_tools=False) + assert {t.category for t in cat} == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS} + assert len(cat) == 6 + + +def test_orchestrator_specs_include_category_field(): + specs = build_tool_specs(orchestrator_catalog()) + by = {s["function"]["name"]: s for s in specs} + assert by["gpt_5_5"]["function"]["category"] == CATEGORY_CLOUD_FRONTIER + assert by["qwen3_5_9b"]["function"]["category"] == CATEGORY_LOCAL_OSS + assert by["web_search"]["function"]["category"] == CATEGORY_BASIC + assert by["shell_exec"]["function"]["category"] == CATEGORY_BASIC + # Every tool carries a category tag. + assert all("category" in s["function"] for s in specs) + + +def test_orchestrator_local_models_get_base_url_when_provided(): + cat = orchestrator_catalog( + local_endpoints={ + "Qwen/Qwen3.5-9B": "http://x/v1", + "Qwen/Qwen3.6-27B-FP8": "http://y/v1", + }) + by = tools_by_name(cat) + assert by["qwen3_5_9b"].base_url == "http://x/v1" + assert by["qwen3_6_27b_fp8"].base_url == "http://y/v1" + # Unmapped local model -> base_url None. + assert by["qwen3_5_122b_a10b_fp8"].base_url is None + # Cloud frontier carries real pricing. + assert by["claude_opus_4_8"].price_in > 0.0 + assert by["gpt_5_5"].price_in > 0.0 + + +def test_openjarvis_tool_bridges_real_tool_with_custom_schema(): + params = { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + } + t = openjarvis_tool("shell_exec", summary="Run a shell command.", params=params) + assert t.kind == KIND_TOOL + assert t.backend_type == "openjarvis-tool" + assert t.model == "shell_exec" + assert t.category == CATEGORY_BASIC + spec = t.to_spec() + assert spec["function"]["name"] == "shell_exec" + assert spec["function"]["parameters"] == params + assert spec["function"]["category"] == CATEGORY_BASIC + + +def test_build_tool_specs_includes_category_for_bridged_tools(): + specs = build_tool_specs([ + openjarvis_tool("calculator", summary="Math.", + params={"type": "object", + "properties": {"expression": {"type": "string"}}, + "required": ["expression"]}), + ]) + assert specs[0]["function"]["category"] == CATEGORY_BASIC + assert "expression" in specs[0]["function"]["parameters"]["properties"] + + def test_to_worker_dict_maps_backend(): cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1") by = tools_by_name(cat) assert to_worker_dict(by["gpt_5"]) == { "name": "gpt_5", "type": "openai", "model": "gpt-5"} - local = to_worker_dict(by["local_model"]) + local = to_worker_dict(by["qwen3_8b"]) assert local["type"] == "vllm" and local["base_url"] == "http://x/v1" diff --git a/tests/test_orchestrator_learning/sft_data/test_datasets.py b/tests/test_orchestrator_learning/sft_data/test_datasets.py new file mode 100644 index 00000000..fe49a782 --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_datasets.py @@ -0,0 +1,136 @@ +"""Offline tests for the reasoning-SFT dataset loaders. + +All loaders take a ``source=`` iterable of raw row dicts so we never touch the +network: we hand them fake rows shaped like the real GeneralThought / +OpenThoughts3 schemas and assert the normalized :class:`Task` fields. +""" + +from __future__ import annotations + +from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import ( + Task, + load_generalthought, + load_openthoughts, +) + +# --- fake rows mirroring the real HF schemas --------------------------------- + +GENERALTHOUGHT_ROWS = [ + { + "question_id": 1, + "question": "What is 2 + 2?", + "reference_answer": "4", + "model_answer": "The answer is 4.", + "task": "High School Math", + "question_source": "Numina/NuminaMath", + }, + { + "question_id": 2, + "question": "How do the neural respiratory centers operate?", + "reference_answer": "Via the medulla oblongata.", + "model_answer": "long answer ...", + "task": "NHS QA", + "question_source": "CogStack/NHSQA", + }, + { + "question_id": 3, + "question": "Sort a list.", + "reference_answer": "", + "model_answer": "use sorted()", + "task": "Sorting", + "question_source": "BAAI/TACO", + }, + # Should be skipped: empty question and answer. + {"question_id": 4, "question": "", "reference_answer": "", "model_answer": ""}, +] + +OPENTHOUGHTS_ROWS = [ + { + "domain": "code", + "source": "stackexchange_codegolf", + "difficulty": 7, + "conversations": [ + {"from": "human", "value": "Write a program that prints hi."}, + {"from": "gpt", "value": "reason Final: print('hi')"}, + ], + }, + { + "domain": "math", + "source": "Numina/NuminaMath", + "difficulty": 5, + "conversations": [ + {"from": "human", "value": "What is 6 times 7?"}, + {"from": "gpt", "value": "6*7=42 The answer is \\boxed{42}."}, + ], + }, + { + "domain": "science", + "source": "sci", + "difficulty": 6, + "conversations": [ + {"from": "human", "value": "What gas do plants release?"}, + {"from": "gpt", "value": "Oxygen."}, + ], + }, + # Should be skipped: no human/gpt turns. + {"domain": "math", "conversations": []}, +] + + +def test_load_generalthought_fields_and_domains(): + tasks = list(load_generalthought(n=10, source=GENERALTHOUGHT_ROWS)) + assert all(isinstance(t, Task) for t in tasks) + # Row 4 (empty) is dropped; remaining 3 yielded. + assert len(tasks) == 3 + by_q = {t.question: t for t in tasks} + assert by_q["What is 2 + 2?"].answer == "4" + assert by_q["What is 2 + 2?"].domain == "math" + assert by_q["What is 2 + 2?"].instruction == "What is 2 + 2?" + assert by_q["How do the neural respiratory centers operate?"].domain == "medical" + # No reference_answer -> falls back to model_answer; TACO source -> code. + assert by_q["Sort a list."].answer == "use sorted()" + assert by_q["Sort a list."].domain == "code" + + +def test_load_generalthought_respects_n(): + tasks = list(load_generalthought(n=2, source=GENERALTHOUGHT_ROWS)) + assert len(tasks) == 2 + + +def test_load_openthoughts_fields_and_domains(): + tasks = list( + load_openthoughts( + n_code=5, n_math=5, n_science=5, source=OPENTHOUGHTS_ROWS + ) + ) + assert all(isinstance(t, Task) for t in tasks) + assert len(tasks) == 3 + by_dom = {t.domain: t for t in tasks} + assert set(by_dom) == {"code", "math", "science"} + # stripped; boxed answer extracted. + assert by_dom["math"].answer == "42" + assert by_dom["math"].question == "What is 6 times 7?" + # Code: think stripped, visible solution kept. + assert "print('hi')" in by_dom["code"].answer + assert by_dom["science"].answer == "Oxygen." + + +def test_load_openthoughts_quota_caps_per_domain(): + rows = [ + { + "domain": "code", + "conversations": [ + {"from": "human", "value": f"q{i}"}, + {"from": "gpt", "value": f"a{i}"}, + ], + } + for i in range(10) + ] + tasks = list(load_openthoughts(n_code=3, n_math=0, n_science=0, source=rows)) + assert len(tasks) == 3 + assert all(t.domain == "code" for t in tasks) + + +def test_task_dataclass_instruction_property(): + t = Task(task_id="x", question="hi?", answer="yes", domain="misc") + assert t.instruction == "hi?" diff --git a/tests/test_orchestrator_learning/sft_data/test_sft_data.py b/tests/test_orchestrator_learning/sft_data/test_sft_data.py deleted file mode 100644 index cc21df7c..00000000 --- a/tests/test_orchestrator_learning/sft_data/test_sft_data.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Offline tests for the orchestrator SFT cold-start pipeline (no network/GPU). - -Covers: ADP transcription -> canonical Episode, difficulty/tier heuristics, -paradigm rendering, reward-ranked selection, serialization round-trip, and the -end-to-end builder via an injected fixture source. -""" - -from __future__ import annotations - -import json -from pathlib import Path - -from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import ( - trajectory_rows_to_episode, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.build import ( - build_sft_dataset, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import ( - PARADIGMS, - render_all, -) -from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best -from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record -from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import ( - Difficulty, - Tier, - covers, - min_covering_tier, - step_difficulty, - tier_telemetry, -) -from openjarvis.learning.intelligence.orchestrator.types import Episode - -# --- fixtures --------------------------------------------------------------- - -_EASY_TURNS = [ - {"source": "user", "class_": "message_action", "content": "What is 2 + 2?"}, - {"source": "agent", "class_": "message_action", "content": "The answer is 4."}, -] - -_CODE_TURNS = [ - { - "source": "user", - "class_": "message_action", - "content": "Sort this list in Python.", - }, - {"source": "agent", "class_": "code_action", "content": "sorted([3, 1, 2])"}, - {"source": "user", "class_": "message_action", "content": "[1, 2, 3]"}, - {"source": "agent", "class_": "message_action", "content": "Sorted: [1, 2, 3]"}, -] - - -def _easy_episode() -> Episode: - ep = trajectory_rows_to_episode("easy-1", _EASY_TURNS) - assert ep is not None - return ep - - -def _code_episode() -> Episode: - ep = trajectory_rows_to_episode("code-1", _CODE_TURNS) - assert ep is not None - return ep - - -# --- adp_loader ------------------------------------------------------------- - - -def test_transcribe_keeps_all_steps_and_problem(): - ep = _code_episode() - assert ep.initial_prompt == "Sort this list in Python." - assert ep.num_turns() == 2 # two agent turns - assert ep.steps[0].action.tool_name == "code" - assert ep.steps[-1].action.is_final_answer is True - assert ep.correct is True # ADP rows are demonstrated solutions - - -def test_transcribe_rejects_empty(): - assert trajectory_rows_to_episode("x", []) is None - assert trajectory_rows_to_episode("x", [{"source": "user", "content": ""}]) is None - - -# --- tiers ------------------------------------------------------------------ - - -def test_difficulty_heuristic(): - assert step_difficulty("message", "short") == Difficulty.EASY - assert step_difficulty("code", "sorted([])") == Difficulty.MED - assert step_difficulty("message", "x", retry_signal=True) == Difficulty.HARD - - -def test_min_covering_tier_and_covers(): - assert min_covering_tier(Difficulty.EASY) == Tier.LOCAL - assert min_covering_tier(Difficulty.MED) == Tier.MID - assert covers(Tier.FRONTIER, Difficulty.HARD) - assert not covers(Tier.LOCAL, Difficulty.HARD) - - -def test_tier_telemetry_local_is_free_cloud_costs(): - assert tier_telemetry("local")["cost_usd"] == 0.0 - assert tier_telemetry("frontier")["cost_usd"] > 0.0 - local_e = tier_telemetry("local")["energy_joules"] - mid_e = tier_telemetry("mid")["energy_joules"] - assert local_e > mid_e - - -# --- paradigms -------------------------------------------------------------- - - -def test_render_all_covers_every_paradigm(): - rendered = render_all(_code_episode()) - assert {r.paradigm for r in rendered} == set(PARADIGMS) - - -def test_baseline_local_incorrect_on_code_task(): - rendered = {r.paradigm: r for r in render_all(_code_episode())} - # code step is MED -> local can't cover it - assert rendered["baseline_local"].predicted_correct is False - assert rendered["baseline_cloud"].predicted_correct is True - assert rendered["toolorchestra"].predicted_correct is True - - -def test_toolorchestra_escalates_code_step(): - rendered = {r.paradigm: r for r in render_all(_code_episode())} - tiers = rendered["toolorchestra"].step_tiers - assert tiers[0] == "mid" # code step escalated off local - assert rendered["baseline_local"].step_tiers[0] == "local" - - -# --- select ----------------------------------------------------------------- - - -def test_select_prefers_local_on_easy_task(): - best = select_best(render_all(_easy_episode())) - assert best is not None - # all-easy -> baseline_local is cheapest-correct - assert best.paradigm == "baseline_local" - - -def test_select_drops_unsolved_task(): - ep = _easy_episode() - # force every paradigm to be predicted-incorrect - renderings = render_all(ep) - for r in renderings: - r.predicted_correct = False - assert select_best(renderings) is None - - -# --- serialize -------------------------------------------------------------- - - -def test_serialize_schema_and_final_answer(): - best = select_best(render_all(_code_episode())) - assert best is not None - rec = to_record(best, reward=0.5) - convo = rec["conversations"] - assert convo[0]["role"] == "system" - assert convo[1]["role"] == "user" - assert any( - c["role"] == "assistant" and "FINAL_ANSWER:" in c["content"] for c in convo - ) - assert rec["paradigm"] == best.paradigm - assert rec["metrics"]["num_steps"] == best.episode.num_turns() - - -# --- build (end-to-end, injected source) ------------------------------------ - - -def test_build_sft_dataset_end_to_end(tmp_path: Path): - def fake_source(*, max_tasks=None, configs=None): - yield _easy_episode() - yield _code_episode() - - out = tmp_path / "traces.jsonl" - stats = build_sft_dataset(str(out), max_tasks=10, source=fake_source) - - assert stats["records_written"] == 2 - assert stats["tasks_seen"] == 2 - lines = out.read_text().strip().splitlines() - assert len(lines) == 2 - for line in lines: - rec = json.loads(line) - assert rec["conversations"][0]["role"] == "system" - - stats_file = out.with_suffix(out.suffix + ".stats.json") - assert stats_file.exists() - assert sum(stats["paradigm_distribution"].values()) == 2 diff --git a/tests/test_orchestrator_learning/sft_data/test_verify.py b/tests/test_orchestrator_learning/sft_data/test_verify.py new file mode 100644 index 00000000..f8a4bdfc --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_verify.py @@ -0,0 +1,107 @@ +"""Offline tests for the answer verifier. + +Math and code paths run with no network. The Gemini judge is only reachable for +non-math/code domains; we monkeypatch ``_gemini_judge`` (or rely on its no-key +fallback) so nothing hits the network. +""" + +from __future__ import annotations + +from openjarvis.learning.intelligence.orchestrator.sft_data import verify as V +from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task +from openjarvis.learning.intelligence.orchestrator.sft_data.verify import ( + make_verifier, + verify_answer, +) + + +def _math(answer: str) -> Task: + return Task(task_id="m", question="q", answer=answer, domain="math") + + +def test_math_boxed_vs_plain_true(): + assert verify_answer(_math("42"), "\\boxed{42}") is True + + +def test_math_boxed_both_sides(): + assert verify_answer(_math("\\boxed{42}"), "The answer is \\boxed{42}") is True + + +def test_math_wrong_number_false(): + assert verify_answer(_math("42"), "41") is False + + +def test_math_numeric_float_match(): + assert verify_answer(_math("42"), "42.0") is True + + +def test_math_symbolic_equivalence(): + # x^2 + 2x + 1 == (x+1)^2 via sympy (skips gracefully if sympy absent). + t = _math("(x+1)^2") + assert verify_answer(t, "x^2 + 2*x + 1") is True + + +def test_empty_prediction_false(): + assert verify_answer(_math("42"), "") is False + + +def test_code_short_gold_substring(): + t = Task(task_id="c", question="print hi", answer="hello", domain="code") + assert verify_answer(t, "the output is hello") is True + + +def test_code_falls_back_to_judge(monkeypatch): + # Long gold + no substring match -> hits judge; monkeypatch to PASS. + monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: True) + t = Task( + task_id="c", + question="solve", + answer="x" * 300, # long -> skips substring shortcut + domain="code", + ) + assert verify_answer(t, "totally different") is True + + +def test_judge_domain_uses_gemini(monkeypatch): + calls = {} + + def fake_judge(task, pred): + calls["hit"] = (task.domain, pred) + return True + + monkeypatch.setattr(V, "_gemini_judge", fake_judge) + t = Task(task_id="s", question="q", answer="Oxygen", domain="science") + assert verify_answer(t, "plants release oxygen") is True + assert calls["hit"][0] == "science" + + +def test_judge_domain_fallback_f1_when_no_key(monkeypatch): + # Judge unavailable (returns None) -> falls back to string / F1. + monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None) + t = Task(task_id="s", question="q", answer="the capital is paris", domain="misc") + assert verify_answer(t, "I think the capital is paris indeed") is True + assert verify_answer(t, "completely unrelated text here") is False + + +def test_make_verifier_reads_final_answer(monkeypatch): + monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None) + verify = make_verifier() + + class Rollout: + final_answer = "\\boxed{42}" + + assert verify(_math("42"), Rollout()) is True + + class BadRollout: + final_answer = "999" + + assert verify(_math("42"), BadRollout()) is False + + +def test_make_verifier_missing_final_answer(): + verify = make_verifier() + + class Empty: + pass + + assert verify(_math("42"), Empty()) is False diff --git a/tests/test_orchestrator_learning/test_build_orchestrator_sft.py b/tests/test_orchestrator_learning/test_build_orchestrator_sft.py new file mode 100644 index 00000000..ad6b9153 --- /dev/null +++ b/tests/test_orchestrator_learning/test_build_orchestrator_sft.py @@ -0,0 +1,124 @@ +"""Offline smoke test for the v1 orchestrator-SFT driver (base self-sampling). + +No network / no GPU: we feed fake ``datasets.Task`` objects, a canned +``UnifiedRollout`` (via a fake ``rollout_fn``), and an accept-all verifier, then +assert ``generate_sft_dataset`` writes a JSONL record in the expected +``conversations`` shape that the SFT trainer consumes. +""" + +from __future__ import annotations + +import json + +from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog +from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn +from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task +from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import ( + generate_sft_dataset, +) + + +def _fake_tasks() -> list[Task]: + return [ + Task(task_id="t-math-1", question="What is 6 * 7?", answer="42", domain="math"), + Task(task_id="t-code-1", question="Reverse 'ab'.", answer="ba", domain="code"), + ] + + +def _canned_rollout(task: Task) -> UnifiedRollout: + # One tool turn + a final-answer turn that echoes the gold answer. + return UnifiedRollout( + turns=[ + UnifiedTurn(reasoning="let me compute", tool_name="code_interpreter", + arguments={"code": "print(6*7)"}, observation="42"), + UnifiedTurn(reasoning="that's the result", tool_name=None), + ], + final_answer=task.answer, cost_usd=0.01, tokens=20, num_tool_calls=1, + ) + + +def test_build_v1_writes_expected_conversations_shape(tmp_path, monkeypatch): + tools = orchestrator_catalog() # specialists unwired (math/coder endpoints None) + + # Accept-all verifier (the real make_verifier may hit Gemini / math checkers). + monkeypatch.setattr( + "openjarvis.learning.intelligence.orchestrator.sft_data.verify.make_verifier", + lambda: (lambda task, rollout: True), + ) + from openjarvis.learning.intelligence.orchestrator.sft_data.verify import ( + make_verifier, + ) + + out = tmp_path / "orchestrator_sft_v1.jsonl" + stats = generate_sft_dataset( + str(out), + tasks=_fake_tasks(), + tools=tools, + rollout_fn=_canned_rollout, + verify_fn=make_verifier(), + samples_per_task=2, + max_keep_per_task=1, + reward_fn=lambda r: -r.cost_usd, + ) + + assert stats["tasks_seen"] == 2 + assert stats["records_written"] == 2 + assert stats["tasks_dropped"] == 0 + + lines = out.read_text().strip().splitlines() + assert len(lines) == 2 + + rec = json.loads(lines[0]) + assert rec["task_id"] == "t-math-1" + assert rec["domain"] == "math" + + roles = [m["role"] for m in rec["conversations"]] + assert roles[0] == "system" and roles[1] == "user" + assert "tool" in roles + assert roles[-1] == "assistant" + # Tool call emitted as a tag the parser reads back. + assert any( + "" in m["content"] and "code_interpreter" in m["content"] + for m in rec["conversations"] + if m["role"] == "assistant" + ) + assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"] + assert (tmp_path / "orchestrator_sft_v1.jsonl.stats.json").exists() + + +def _load_driver(): + """Load the (non-package) CLI driver module by path.""" + import importlib.util + from pathlib import Path + + root = Path(__file__).resolve().parents[2] + path = root / "scripts" / "orchestrator" / "build_orchestrator_sft.py" + spec = importlib.util.spec_from_file_location("build_orchestrator_sft", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_driver_runs_end_to_end_with_fakes(tmp_path, monkeypatch): + """Drive ``main`` with no network: fake task loader, rollout, and verifier.""" + drv = _load_driver() + + monkeypatch.setattr(drv, "load_sft_tasks", _fake_tasks) + monkeypatch.setattr(drv, "run_unified_rollout", + lambda question, tools, **kw: _canned_rollout(_fake_tasks()[0])) + monkeypatch.setattr(drv, "make_verifier", lambda: (lambda task, rollout: True)) + # make_call_orchestrator builds an OpenAI client lazily, so it never connects + # here (run_unified_rollout is stubbed out). + + out = tmp_path / "v1.jsonl" + rc = drv.main([ + "--out", str(out), + "--samples-per-task", "1", + "--max-tasks", "2", + ]) + assert rc == 0 + lines = out.read_text().strip().splitlines() + assert len(lines) == 2 + rec = json.loads(lines[0]) + assert rec["conversations"][0]["role"] == "system" + assert "FINAL_ANSWER" in rec["conversations"][-1]["content"] diff --git a/tests/test_orchestrator_learning/test_eval_backend.py b/tests/test_orchestrator_learning/test_eval_backend.py new file mode 100644 index 00000000..2d09133e --- /dev/null +++ b/tests/test_orchestrator_learning/test_eval_backend.py @@ -0,0 +1,170 @@ +"""Offline tests for the orchestrator eval backend + the eval CLI's +benchmark-name normalization. + +No network / models: ``run_unified_rollout`` is monkeypatched to return a +canned rollout, and the dataset-key mapping is asserted against the real +``openjarvis.evals.cli`` registry without instantiating any (network-bound) +dataset. +""" + +from __future__ import annotations + +import pytest + +from openjarvis.learning.intelligence.orchestrator import eval_backend as eb +from openjarvis.learning.intelligence.orchestrator.eval_backend import ( + OrchestratorBackend, +) + + +class _CannedRollout: + """Minimal stand-in for UnifiedRollout.""" + + def __init__(self, final_answer: str) -> None: + self.final_answer = final_answer + self.cost_usd = 0.123 + self.tokens = 42 + self.num_tool_calls = 1 + self.parse_failures = 0 + self.turns = [object()] + + def tool_calls(self): + return [("web_search", {"query": "x"})] + + +def _patch_rollout(monkeypatch, rollout): + """Patch run_unified_rollout where the backend looks it up.""" + monkeypatch.setattr( + eb._rollout_mod, + "run_unified_rollout", + lambda *a, **k: rollout, + ) + + +def test_construct_is_network_free(): + """Constructing the backend builds the catalog but touches no network.""" + backend = OrchestratorBackend() + assert backend.backend_id == "orchestrator" + assert backend.orchestrator_endpoint == eb.DEFAULT_ENDPOINT + assert backend.orchestrator_model == eb.DEFAULT_MODEL + # Catalog is non-empty (cloud frontier + local OSS + basic tools). + assert len(backend._tools) > 0 + names = {t.name for t in backend._tools} + assert "web_search" in names + + +def test_local_endpoints_mapping(): + """local_endpoints (model-id -> base_url) wires the local OSS tool.""" + backend = OrchestratorBackend( + local_endpoints={"Qwen/Qwen3.5-9B": "http://m:9/v1"}, + ) + assert backend.local_endpoints["Qwen/Qwen3.5-9B"] == "http://m:9/v1" + # The corresponding local-OSS tool should carry that base_url. + tool = next(t for t in backend._tools if t.model == "Qwen/Qwen3.5-9B") + assert tool.base_url == "http://m:9/v1" + + +def test_generate_returns_final_answer(monkeypatch): + """generate() returns the rollout's final_answer string.""" + _patch_rollout(monkeypatch, _CannedRollout("42")) + backend = OrchestratorBackend() + out = backend.generate("what is 6*7?", model="qwen3-8b") + assert out == "42" + + +def test_generate_full_payload(monkeypatch): + """generate_full() returns the runner-expected payload shape.""" + _patch_rollout(monkeypatch, _CannedRollout("42")) + backend = OrchestratorBackend() + full = backend.generate_full("q", model="qwen3-8b") + assert full["content"] == "42" + assert full["cost_usd"] == 0.123 + assert full["usage"]["completion_tokens"] == 42 + assert full["tool_calls"] == 1 + assert full["turn_count"] == 1 + assert full["framework"] == "openjarvis-orchestrator" + assert "error" not in full + assert full["latency_seconds"] >= 0.0 + + +def test_generate_full_error_path(monkeypatch): + """An exception inside the rollout becomes a recorded error, not a crash.""" + + def _boom(*a, **k): + raise RuntimeError("vllm down") + + monkeypatch.setattr(eb._rollout_mod, "run_unified_rollout", _boom) + backend = OrchestratorBackend() + full = backend.generate_full("q", model="qwen3-8b") + assert full["content"] == "" + assert "vllm down" in full["error"] + assert full["usage"]["completion_tokens"] == 0 + + +# --------------------------------------------------------------------------- +# Eval-script benchmark name mapping +# --------------------------------------------------------------------------- + +# The 5 benchmarks the eval script targets, by the script's (alias) names. +SCRIPT_BENCHMARKS = [ + "gaia", + "terminalbench_v2_1", + "taubench", + "mmlu_pro", + "supergpqa", +] + + +def _load_script(): + import importlib.util + from pathlib import Path + + repo_root = Path(__file__).resolve().parents[2] + script_path = repo_root / "scripts" / "orchestrator" / "eval_orchestrator.py" + spec = importlib.util.spec_from_file_location("eval_orchestrator", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_benchmark_aliases_map_to_registry_keys(): + """Each script benchmark normalizes to a real evals.cli BENCHMARKS key.""" + from openjarvis.evals.cli import BENCHMARKS + + script = _load_script() + for name in SCRIPT_BENCHMARKS: + key = script._normalize_benchmark(name) + assert key in BENCHMARKS, f"{name} -> {key} not in BENCHMARKS registry" + + +def test_normalize_specific_keys(): + """Spot-check the underscored aliases normalize to the dotted/dashed keys.""" + script = _load_script() + assert script._normalize_benchmark("terminalbench_v2_1") == "terminalbench-v2.1" + assert script._normalize_benchmark("mmlu_pro") == "mmlu-pro" + assert script._normalize_benchmark("gaia") == "gaia" + + +def test_build_dataset_accepts_keys(): + """_build_dataset must recognise the normalized keys. + + We only assert it doesn't raise the 'Unknown benchmark' ClickException; + datasets that need network/files at construction are xfail-skipped. + """ + import click + + from openjarvis.evals.cli import _build_dataset + + script = _load_script() + for name in SCRIPT_BENCHMARKS: + key = script._normalize_benchmark(name) + try: + _build_dataset(key) + except click.ClickException as exc: + if "Unknown benchmark" in str(exc): + pytest.fail(f"{key} not recognised by _build_dataset") + # Other ClickExceptions (missing files/creds) are acceptable here. + except Exception: + # Construction may require network/files — that's fine; the point + # is the key is recognised (didn't hit the unknown-benchmark else). + pytest.skip(f"{key} dataset construction needs resources") diff --git a/tests/test_orchestrator_learning/test_grpo_epoch.py b/tests/test_orchestrator_learning/test_grpo_epoch.py new file mode 100644 index 00000000..5afaef0d --- /dev/null +++ b/tests/test_orchestrator_learning/test_grpo_epoch.py @@ -0,0 +1,408 @@ +"""Offline (no-GPU) tests for OrchestratorGRPOTrainer._train_epoch. + +We never touch torch/model/network: ``load_grpo_prompts`` is monkeypatched +to return a handful of fake Tasks, and ``_grpo_step`` is monkeypatched to +return canned metrics. The test asserts ``_train_epoch`` chunks the prompt +pool correctly and aggregates per-chunk metrics. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +import openjarvis.learning.intelligence.orchestrator.policy_model as policy_model_mod +import openjarvis.learning.intelligence.orchestrator.sft_data.datasets as datasets +from openjarvis.learning.intelligence.orchestrator.grpo_trainer import ( + OrchestratorGRPOConfig, + OrchestratorGRPOTrainer, + RolloutRecord, +) +from openjarvis.learning.intelligence.orchestrator.policy_model import ( + OrchestratorPolicyModel, +) + + +@dataclass +class FakeTask: + task_id: str + question: str + answer: str + domain: str = "math" + + @property + def instruction(self) -> str: + return self.question + + +def _make_tasks(n: int): + return [ + FakeTask(task_id=f"t{i}", question=f"q{i}", answer=str(i)) + for i in range(n) + ] + + +def _build_trainer(monkeypatch, n_prompts: int, prompts_per_step: int): + """Build a trainer with no real model and patched data loading. + + ``from_pretrained`` is stubbed to a model-less policy so construction + touches no GPU / network even when torch happens to be installed. + """ + tasks = _make_tasks(n_prompts) + monkeypatch.setattr(datasets, "load_grpo_prompts", lambda *a, **k: tasks) + monkeypatch.setattr( + policy_model_mod.OrchestratorPolicyModel, + "from_pretrained", + classmethod(lambda cls, *a, **k: OrchestratorPolicyModel(model=None)), + ) + + cfg = OrchestratorGRPOConfig( + grpo_max_prompts=n_prompts, + prompts_per_step=prompts_per_step, + num_samples_per_prompt=2, + ) + trainer = OrchestratorGRPOTrainer(cfg) + return trainer, tasks + + +def test_train_epoch_chunks_and_aggregates(monkeypatch): + trainer, tasks = _build_trainer(monkeypatch, n_prompts=10, prompts_per_step=4) + + calls = [] + + def fake_step(chunk): + calls.append(len(chunk)) + return { + "loss": 1.0, + "reward": 0.5, + "accuracy": 1.0, + "n_prompts": len(chunk), + } + + monkeypatch.setattr(trainer, "_grpo_step", fake_step) + + metrics = trainer._train_epoch(epoch=0) + + # 10 prompts / 4 per step => chunks of [4, 4, 2] + assert calls == [4, 4, 2] + assert metrics["n_batches"] == 3 + assert metrics["n_prompts"] == 10 + # Weighted means: all chunks report the same per-prompt values. + assert abs(metrics["loss"] - 1.0) < 1e-9 + assert abs(metrics["reward"] - 0.5) < 1e-9 + assert abs(metrics["accuracy"] - 1.0) < 1e-9 + assert metrics["epoch"] == 0 + + +def test_train_epoch_exact_division(monkeypatch): + trainer, _ = _build_trainer(monkeypatch, n_prompts=8, prompts_per_step=4) + chunks = [] + monkeypatch.setattr( + trainer, + "_grpo_step", + lambda chunk: ( + chunks.append(len(chunk)), + {"loss": 0.0, "reward": 0.0, "accuracy": 0.0, "n_prompts": len(chunk)}, + )[1], + ) + metrics = trainer._train_epoch(epoch=2) + assert chunks == [4, 4] + assert metrics["n_batches"] == 2 + assert metrics["n_prompts"] == 8 + assert metrics["epoch"] == 2 + + +def test_train_epoch_weighted_mean(monkeypatch): + trainer, _ = _build_trainer(monkeypatch, n_prompts=6, prompts_per_step=4) + + # First chunk (4 prompts) reward=1.0, second chunk (2 prompts) reward=0.0 + seq = iter([1.0, 0.0]) + + def fake_step(chunk): + return { + "loss": 0.0, + "reward": next(seq), + "accuracy": 0.0, + "n_prompts": len(chunk), + } + + monkeypatch.setattr(trainer, "_grpo_step", fake_step) + metrics = trainer._train_epoch(epoch=0) + # weighted: (1.0*4 + 0.0*2) / 6 = 0.6667 + assert abs(metrics["reward"] - (4.0 / 6.0)) < 1e-9 + + +def test_prompts_cached(monkeypatch): + trainer, _ = _build_trainer(monkeypatch, n_prompts=4, prompts_per_step=4) + n_loads = {"count": 0} + + orig = datasets.load_grpo_prompts + + def counting(*a, **k): + n_loads["count"] += 1 + return orig(*a, **k) + + monkeypatch.setattr(datasets, "load_grpo_prompts", counting) + monkeypatch.setattr( + trainer, + "_grpo_step", + lambda chunk: {"loss": 0.0, "reward": 0.0, "accuracy": 0.0, "n_prompts": len(chunk)}, + ) + trainer._train_epoch(epoch=0) + trainer._train_epoch(epoch=1) + # Loaded once, cached on the trainer thereafter. + assert n_loads["count"] == 1 + + +def test_train_epoch_no_model_still_iterates(monkeypatch): + """Even with policy.model None (no GPU), epoch iterates via _grpo_step.""" + trainer, _ = _build_trainer(monkeypatch, n_prompts=5, prompts_per_step=2) + assert trainer.policy.model is None # CPU/no-torch stub + called = {"n": 0} + + def fake_step(chunk): + called["n"] += 1 + return {"loss": 0.0, "reward": 0.0, "accuracy": 0.0, "n_prompts": len(chunk)} + + monkeypatch.setattr(trainer, "_grpo_step", fake_step) + metrics = trainer._train_epoch(epoch=0) + assert called["n"] == 3 # [2, 2, 1] + assert metrics["n_prompts"] == 5 + + +# --------------------------------------------------------------------------- +# _grpo_step over the unified rollout (offline, fake model + patched paths) +# --------------------------------------------------------------------------- + + +torch = pytest.importorskip("torch") + + +class _FakeModel: + """Stand-in for the HF policy/ref model: trainable enough to backprop. + + ``parameters()`` returns no params so the NaN-grad scan and + ``clip_grad_norm_`` are no-ops; ``train()`` is a no-op. The real + log-prob / generation paths are monkeypatched, so the model is never + actually called. + """ + + def parameters(self): + return iter(()) + + def train(self): + return self + + +class _FakeOptimizer: + def zero_grad(self): + pass + + def step(self): + pass + + +def _build_step_trainer(monkeypatch, *, group_size=4): + """Trainer wired for a CPU ``_grpo_step`` run: fake model + optimizer, + device='cpu', and the rollout/logprob paths left for the test to patch. + """ + tasks = _make_tasks(3) + monkeypatch.setattr(datasets, "load_grpo_prompts", lambda *a, **k: tasks) + monkeypatch.setattr( + policy_model_mod.OrchestratorPolicyModel, + "from_pretrained", + classmethod(lambda cls, *a, **k: OrchestratorPolicyModel(model=None)), + ) + cfg = OrchestratorGRPOConfig( + grpo_max_prompts=3, + prompts_per_step=3, + num_samples_per_prompt=group_size, + ) + trainer = OrchestratorGRPOTrainer(cfg) + # Make it look like a real (CPU) training setup. + trainer.policy.model = _FakeModel() + trainer.ref_policy.model = _FakeModel() + trainer.optimizer = _FakeOptimizer() + trainer.device = torch.device("cpu") + return trainer, tasks + + +def test_grpo_step_uses_unified_rollout_not_old_env(monkeypatch): + """_grpo_step samples groups via the policy rollout caller, normalises + advantages within the group, and never touches the old + THOUGHT/TOOL/INPUT environment path.""" + + # Tripwire: the old environment generation path must NOT be invoked. + import openjarvis.learning.intelligence.orchestrator.policy_model as pm + + def _boom(*a, **k): # pragma: no cover - only fires on regression + raise AssertionError( + "GRPO generation hit the old THOUGHT/TOOL/INPUT path " + "(policy_model._parse_output / predict_action)" + ) + + monkeypatch.setattr(pm.OrchestratorPolicyModel, "_parse_output", _boom) + monkeypatch.setattr(pm.OrchestratorPolicyModel, "predict_action", _boom) + monkeypatch.setattr(pm.OrchestratorPolicyModel, "_build_prompt", _boom) + + trainer, tasks = _build_step_trainer(monkeypatch, group_size=4) + + # verify_answer: even-indexed rollouts correct, odd incorrect -> spread of + # rewards inside each group so std > 0 and advantages get normalised. + import openjarvis.learning.intelligence.orchestrator.sft_data.verify as verify_mod + + seen_answers = [] + + def fake_verify(task, pred): + seen_answers.append(pred) + return pred.endswith("correct") + + monkeypatch.setattr(verify_mod, "verify_answer", fake_verify) + + def fake_rollout_group(task): + recs = [] + for i in range(trainer.config.num_samples_per_prompt): + ans = f"ans-{i}-correct" if i % 2 == 0 else f"ans-{i}-wrong" + recs.append( + RolloutRecord( + final_answer=ans, + cost_usd=0.01 * (i + 1), + token_ids=[1, 2, 3, 4], + assistant_mask=[0, 1, 0, 1], + ) + ) + return recs + + monkeypatch.setattr(trainer, "_rollout_group", fake_rollout_group) + + def fake_logprob(token_ids, assistant_mask, *, ref): + # token_ids/mask flow through unchanged from the records. + assert token_ids == [1, 2, 3, 4] + assert assistant_mask == [0, 1, 0, 1] + # Distinct, grad-bearing values for current vs ref so the ratio is + # finite and backward() builds a real graph. + base = -1.0 if ref else -0.9 + return torch.tensor(base, requires_grad=not ref) + + monkeypatch.setattr(trainer, "_trajectory_logprob", fake_logprob) + + metrics = trainer._grpo_step(tasks) + + # Metrics aggregate over the 3 tasks. + assert metrics["n_prompts"] == 3 + # Half of each group is correct -> accuracy 0.5. + assert abs(metrics["accuracy"] - 0.5) < 1e-9 + assert "loss" in metrics and "reward" in metrics + # The unified-rollout caller produced our canned answers (proves we went + # through the rollout path, not the old env). + assert any(a.endswith("correct") for a in seen_answers) + assert any(a.endswith("wrong") for a in seen_answers) + + +def test_grpo_step_group_normalised_advantages(monkeypatch): + """Advantages within a group are mean-0 / std-1 (group-relative).""" + trainer, tasks = _build_step_trainer(monkeypatch, group_size=4) + + import openjarvis.learning.intelligence.orchestrator.sft_data.verify as verify_mod + + monkeypatch.setattr( + verify_mod, "verify_answer", lambda task, pred: pred == "good" + ) + + all_groups = [] + + def fake_rollout_group(task): + # 2 correct, 2 wrong, varied costs -> non-degenerate reward spread. + answers = ["good", "good", "bad", "bad"] + recs = [ + RolloutRecord( + final_answer=answers[i], + cost_usd=0.0, + token_ids=[5, 6, 7], + assistant_mask=[0, 1, 1], + ) + for i in range(4) + ] + all_groups.append(recs) + return recs + + monkeypatch.setattr(trainer, "_rollout_group", fake_rollout_group) + + def fake_logprob(token_ids, assistant_mask, *, ref): + return torch.tensor(0.0, requires_grad=not ref) + + monkeypatch.setattr(trainer, "_trajectory_logprob", fake_logprob) + + trainer._grpo_step(tasks) + + # 3 tasks, each a group of 4. Read the advantages _grpo_step stamped onto + # each record; every group must be group-normalised (mean ~0, std ~1). + assert len(all_groups) == 3 + for recs in all_groups: + group = [r.advantage for r in recs] + assert len(group) == 4 + mean = sum(group) / 4 + assert abs(mean) < 1e-6 + var = sum((x - mean) ** 2 for x in group) / 4 + assert abs(var - 1.0) < 1e-6 + + +def test_stitch_trajectory_assistant_mask(): + """Stitching per-turn (prompt_ids, gen_ids) yields a correct mask: 1 only + on the policy-generated tokens, growing prompts diffed across turns.""" + trainer = OrchestratorGRPOTrainer.__new__(OrchestratorGRPOTrainer) + # Turn 1: prompt [10,11], gen [20,21]. Turn 2 prompt is the full prior + # transcript [10,11,20,21] + a new observation token [12], gen [22]. + turn_buffer = [ + ([10, 11], [20, 21]), + ([10, 11, 20, 21, 12], [22]), + ] + token_ids, mask = trainer._stitch_trajectory(turn_buffer) + # Turn1: new prompt [10,11] mask0, gen [20,21] mask1. + # Turn2: new prompt [12] mask0 (the [10,11,20,21] already emitted), gen [22] mask1. + assert token_ids == [10, 11, 20, 21, 12, 22] + assert mask == [0, 0, 1, 1, 0, 1] + + +def test_trajectory_logprob_masks_to_assistant_positions(monkeypatch): + """_trajectory_logprob sums log-probs only over assistant-masked tokens, + using logits at position i-1 to predict token i.""" + trainer = OrchestratorGRPOTrainer.__new__(OrchestratorGRPOTrainer) + trainer.device = torch.device("cpu") + + token_ids = [3, 7, 5, 9] + assistant_mask = [0, 1, 0, 1] # tokens at idx 1 and 3 are assistant + + vocab = 10 + seq = len(token_ids) + # Deterministic logits so the expected sum is computable. + logits = torch.zeros(1, seq, vocab) + for i in range(seq): + logits[0, i, :] = torch.arange(vocab, dtype=torch.float32) * 0.1 * (i + 1) + + class _LogitModel: + def __call__(self, input_ids=None): + class _Out: + pass + + o = _Out() + o.logits = logits + return o + + def parameters(self): + return iter(()) + + trainer.policy = OrchestratorPolicyModel(model=_LogitModel()) + trainer.ref_policy = OrchestratorPolicyModel(model=_LogitModel()) + + out = trainer._trajectory_logprob(token_ids, assistant_mask, ref=True) + + # Expected: log_softmax(logits[i-1])[token_i] summed over masked i in {1,3}. + expected = 0.0 + for i in range(1, seq): + if assistant_mask[i] != 1: + continue + lp = torch.log_softmax(logits[0, i - 1, :], dim=-1)[token_ids[i]] + expected += lp.item() + assert abs(out.item() - expected) < 1e-5 diff --git a/tests/test_orchestrator_learning/test_reward_costaware.py b/tests/test_orchestrator_learning/test_reward_costaware.py new file mode 100644 index 00000000..0a38d723 --- /dev/null +++ b/tests/test_orchestrator_learning/test_reward_costaware.py @@ -0,0 +1,86 @@ +"""Tests for the cost-aware GRPO reward (offline, no GPU).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from openjarvis.learning.intelligence.orchestrator.reward import ( + CostAwareReward, + lambda_sweep, +) + + +@dataclass +class FakeEpisode: + """Minimal stand-in exposing the two fields CostAwareReward reads.""" + + correct: bool + total_cost_usd: float = 0.0 + + +class TestCostAwareReward: + def test_correct_zero_cost_is_plus_one(self): + r = CostAwareReward(lam=0.0) + assert r.compute(FakeEpisode(correct=True, total_cost_usd=0.0)) == 1.0 + + def test_incorrect_zero_cost_is_minus_one(self): + r = CostAwareReward(lam=0.0) + assert r.compute(FakeEpisode(correct=False, total_cost_usd=0.0)) == -1.0 + + def test_correct_with_lam_penalty_when_cost(self): + # base=+1, penalty = lam * cost / cost_max = 0.2 * 0.05 / 0.10 = 0.1 + r = CostAwareReward(lam=0.2, cost_max=0.10) + val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.05)) + assert abs(val - (1.0 - 0.1)) < 1e-9 + + def test_lam_zero_ignores_cost(self): + r = CostAwareReward(lam=0.0, cost_max=0.10) + val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.5)) + assert val == 1.0 + + def test_cost_at_budget_full_penalty(self): + # cost == cost_max → penalty == lam exactly. + r = CostAwareReward(lam=0.4, cost_max=0.10) + val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.10)) + assert abs(val - (1.0 - 0.4)) < 1e-9 + + def test_incorrect_also_penalised(self): + r = CostAwareReward(lam=0.2, cost_max=0.10) + val = r.compute(FakeEpisode(correct=False, total_cost_usd=0.05)) + assert abs(val - (-1.0 - 0.1)) < 1e-9 + + def test_breakdown_fields(self): + r = CostAwareReward(lam=0.2, cost_max=0.10) + b = r.compute_with_breakdown(FakeEpisode(correct=True, total_cost_usd=0.05)) + assert set(b) == {"correct", "base", "cost_term", "reward"} + assert b["correct"] == 1.0 + assert b["base"] == 1.0 + assert abs(b["cost_term"] - (-0.1)) < 1e-9 + assert abs(b["reward"] - 0.9) < 1e-9 + + def test_breakdown_reward_matches_compute(self): + r = CostAwareReward(lam=0.3, cost_max=0.10) + ep = FakeEpisode(correct=False, total_cost_usd=0.07) + assert abs(r.compute_with_breakdown(ep)["reward"] - r.compute(ep)) < 1e-12 + + def test_compute_batch(self): + r = CostAwareReward(lam=0.0) + eps = [FakeEpisode(correct=True), FakeEpisode(correct=False)] + assert r.compute_batch(eps) == [1.0, -1.0] + + +class TestLambdaSweep: + def test_sweep_length(self): + values = [0.0, 0.05, 0.1, 0.2, 0.4] + rewards = lambda_sweep(values) + assert len(rewards) == len(values) + + def test_sweep_assigns_lambdas(self): + values = [0.0, 0.05, 0.1, 0.2, 0.4] + rewards = lambda_sweep(values, cost_max=0.25) + assert [r.lam for r in rewards] == values + assert all(r.cost_max == 0.25 for r in rewards) + + def test_sweep_returns_costaware(self): + rewards = lambda_sweep([0.1]) + assert isinstance(rewards[0], CostAwareReward) diff --git a/tests/test_orchestrator_learning/test_sft_trainer.py b/tests/test_orchestrator_learning/test_sft_trainer.py index 941427f3..f05ab60d 100644 --- a/tests/test_orchestrator_learning/test_sft_trainer.py +++ b/tests/test_orchestrator_learning/test_sft_trainer.py @@ -15,7 +15,7 @@ from openjarvis.learning.intelligence.orchestrator.sft_trainer import ( class TestOrchestratorSFTConfig: def test_defaults(self): cfg = OrchestratorSFTConfig() - assert cfg.model_name == "Qwen/Qwen3-1.7B" + assert cfg.model_name == "Qwen/Qwen3.5-9B" assert cfg.num_epochs == 3 assert cfg.batch_size == 8 assert cfg.learning_rate == 2e-5