"""Bounded online agent GRPO through pinned TRL native tool environments.
Only adapter configuration, raw-token preservation and audit logic live here.
Generation, reference log probabilities, policy loss and optimization stay in TRL.
"""
from __future__ import annotations
import argparse
import copy
from datetime import datetime, timezone
import hashlib
import importlib.metadata
import inspect
import json
import math
import os
from pathlib import Path
import statistics
from typing import Any
from .evaluation import checkpoint_sha256
from .training import (_keys, _positive, _text, chat_template_kwargs, read_json_snapshot,
read_jsonl_snapshot, require_versions, sha256, write_manifest)
DEFAULT_TRAINING = {
"max_steps": 32, "num_generations": 4, "per_device_train_batch_size": 4,
"gradient_accumulation_steps": 1, "max_completion_length": 768,
"max_tool_calling_iterations": 3, "learning_rate": 1e-5, "beta": 0.02,
"temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0.0,
}
DEFAULT_LORA = {"r": 16, "alpha": 32, "dropout": 0.0, "target_modules": "all-linear"}
FORMAT_VERSION = "qwen3-sampled-assistant-ids-v2"
REWARD_PROTOCOL = {
"version": "trusted-milestones-v1",
"scope": "training_only",
"formula": "clip(0.15*lookup_success + 0.25*calculation_success + 0.60*strict_success - 0.05*trusted_tool_errors, 0, 1)",
"milestones": "one-time authoritative environment flags; never assistant claims or repeated-call counts",
"strict_success_metric": "unchanged binary score_episode.success and score_episode.reward",
"limitation": "Partial workflows can earn up to 0.40 without solving the task; mean training reward is not a success rate.",
}
def milestone_training_reward(score: dict[str, Any]) -> dict[str, Any]:
"""Give training credit only for independently verified environment progress.
The original strict score remains unchanged for evaluation and reporting.
Repeated valid calls cannot increase a one-time milestone. Only actual
failed tool executions in the trusted environment trace incur penalties.
"""
for key in ("lookup_success", "calculation_success", "success"):
if type(score.get(key)) is not bool:
raise ValueError("Milestone reward requires authoritative boolean environment flags")
trace = score.get("trace")
if not isinstance(trace, list) or any(not isinstance(row, dict) or type(row.get("ok")) is not bool for row in trace):
raise ValueError("Milestone reward requires the trusted tool trace with boolean ok fields")
tool_errors = sum(not row["ok"] for row in trace)
if score["calculation_success"] and not score["lookup_success"]:
raise ValueError("A verified calculation must follow a verified lookup")
if score["success"] and (not score["calculation_success"] or tool_errors):
raise ValueError("Strict success cannot coexist with missing milestones or tool errors")
strict_reward = score.get("reward")
if type(strict_reward) not in (int, float) or strict_reward != float(score["success"]):
raise ValueError("Environment score.reward must remain the unchanged binary success indicator")
lookup_bonus = 0.15 * score["lookup_success"]
calculation_bonus = 0.25 * score["calculation_success"]
strict_bonus = 0.60 * score["success"]
penalty = 0.05 * tool_errors
reward = max(0.0, min(1.0, lookup_bonus + calculation_bonus + strict_bonus - penalty))
return {"version": REWARD_PROTOCOL["version"], "training_reward": reward,
"lookup_bonus": lookup_bonus, "calculation_bonus": calculation_bonus,
"strict_success_bonus": strict_bonus, "trusted_tool_errors": tool_errors,
"tool_error_penalty": penalty, "strict_success": score["success"],
"strict_reward": float(strict_reward)}
def _raw_preserving_template(native: str) -> str:
"""Preserve generated assistant bytes; render only external tool framing anew."""
start = ' {%- elif message.role == "assistant" %}\n'
end = " {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}"
if native.count(start) != 1 or native.count(end) != 1:
raise ValueError("Pinned TRL Qwen template layout changed; review raw-span formatting")
raw = (
start + ' {%- if message.raw_response is defined %}\n'
" {{- '<|im_start|>assistant\\n' }}\n"
' {%- if enable_thinking is defined and enable_thinking is false %}\n'
" {{- '\\n\\n\\n\\n' }}\n"
' {%- endif %}\n'
" {{- message.hessian_raw_span_start | default('') }}\n"
' {{- message.raw_response }}\n'
" {{- message.hessian_raw_span_end | default('') }}\n"
" {%- if not message.raw_response.endswith('<|im_end|>') %}{{- '<|im_end|>' }}{%- endif %}\n"
" {{- '\\n' }}\n"
' {%- else %}\n'
)
return native.replace(start, raw).replace(
end, " {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- elif message.role == \"tool\" %}")
def _install_token_preserving_chat_template(tokenizer: Any) -> None:
"""Keep sampled assistant IDs through native TRL's repeated history rendering.
The native tool loop renders once for its length preflight and again for
generation. Both must see the same exact token history. Marker fields exist
only in private rendering copies; they never enter text returned to callers
or model inputs. Only external text is encoded again.
"""
if getattr(tokenizer, "_hessian_exact_template_configured", False):
return
original_apply = tokenizer.apply_chat_template
def apply_exact(conversation: Any, *args: Any, **kwargs: Any) -> Any:
if not kwargs.get("tokenize", True):
return original_apply(conversation, *args, **kwargs)
batched = bool(conversation) and isinstance(conversation[0], (list, tuple))
conversations = list(conversation) if batched else [conversation]
assistant_messages = [
message for messages in conversations for message in messages
if isinstance(message, dict) and message.get("role") == "assistant"
and ("raw_response" in message or "raw_token_ids" in message)
]
if not assistant_messages:
return original_apply(conversation, *args, **kwargs)
if args:
raise ValueError("Exact sampled-token history requires named rendering options")
if (kwargs.get("truncation", False) not in (False, None)
or kwargs.get("continue_final_message", False)
or kwargs.get("return_assistant_tokens_mask", False)
or kwargs.get("tokenizer_kwargs") not in (None, {})):
raise ValueError("Unsupported tokenization option for exact sampled-token history")
supported_returns = {"return_tensors", "return_dict", "return_assistant_tokens_mask"}
if any(key.startswith("return_") and key not in supported_returns for key in kwargs):
raise ValueError("Unsupported output option for exact sampled-token history")
tensors = kwargs.get("return_tensors")
if tensors not in (None, "pt", "np"):
raise ValueError("Unsupported tensors for exact sampled-token history")
marker_keys = {"hessian_raw_span_start", "hessian_raw_span_end"}
if any(marker_keys.intersection(message) for messages in conversations
for message in messages if isinstance(message, dict)):
raise ValueError("Reserved sampled-token marker fields in conversation")
for message in assistant_messages:
ids = message.get("raw_token_ids")
if (not isinstance(ids, list)
or any(type(token) is not int or token < 0 for token in ids)
or not isinstance(message.get("raw_response"), str)):
raise ValueError("Assistant history requires captured raw token IDs and text")
if tokenizer.decode(ids, skip_special_tokens=False,
clean_up_tokenization_spaces=False) != message["raw_response"]:
raise ValueError("Captured assistant token IDs differ from raw response text")
render_options = dict(kwargs)
for key in ("tokenize", "padding", "padding_side", "truncation", "max_length",
"return_tensors", "return_dict", "return_assistant_tokens_mask",
"tokenizer_kwargs"):
render_options.pop(key, None)
unmarked = original_apply(conversations, tokenize=False, **render_options)
if not isinstance(unmarked, list) or len(unmarked) != len(conversations):
raise ValueError("Unexpected native batch rendering shape")
# Deterministic collision avoidance also covers tool observations and
# template text, rather than assuming sampled text cannot contain markers.
counter = 0
marker_prefix = "\x00hessian-sampled-span:0:"
while any(marker_prefix in text for text in unmarked):
counter += 1
marker_prefix = "\x00hessian-sampled-span:" + str(counter) + ":"
marked = [[copy.deepcopy(message) for message in messages] for messages in conversations]
span_rows = []
span_index = 0
for messages in marked:
spans = []
for message in messages:
if message.get("role") != "assistant" or "raw_token_ids" not in message:
continue
start = marker_prefix + str(span_index) + ":start\x00"
end = marker_prefix + str(span_index) + ":end\x00"
message["hessian_raw_span_start"] = start
message["hessian_raw_span_end"] = end
spans.append((start, end, message["raw_token_ids"], message["raw_response"]))
span_index += 1
span_rows.append(spans)
rendered = original_apply(marked, tokenize=False, **render_options)
if not isinstance(rendered, list) or len(rendered) != len(conversations):
raise ValueError("Unexpected marked native batch rendering shape")
rows = []
for text, expected, spans in zip(rendered, unmarked, span_rows, strict=True):
ids, cursor, clean_parts = [], 0, []
for start, end, sampled_ids, raw in spans:
if text.count(start) != 1 or text.count(end) != 1:
raise ValueError("Native template omitted or duplicated sampled-token markers")
left, right = text.index(start), text.index(end)
if left < cursor or right < left + len(start):
raise ValueError("Native template reordered sampled-token spans")
external = text[cursor:left]
if text[left + len(start):right] != raw:
raise ValueError("Native template changed captured assistant text")
ids.extend(tokenizer.encode(external, add_special_tokens=False))
ids.extend(sampled_ids)
clean_parts.extend((external, raw))
cursor = right + len(end)
external = text[cursor:]
ids.extend(tokenizer.encode(external, add_special_tokens=False))
clean_parts.append(external)
if ("".join(clean_parts) != expected
or tokenizer.decode(ids, skip_special_tokens=False,
clean_up_tokenization_spaces=False) != expected):
raise ValueError("Exact sampled-token history changed native rendered text")
rows.append(ids)
# The native first turn has no captured history and uses its original
# tokenizer path. Later calls retain the same batch/tensor/padding shapes.
inputs = rows if batched or tensors is not None else rows[0]
encoded = tokenizer.pad(
{"input_ids": inputs}, padding=kwargs.get("padding", False),
max_length=kwargs.get("max_length"),
padding_side=kwargs.get("padding_side", tokenizer.padding_side),
return_attention_mask=True, return_tensors=tensors,
)
return encoded if kwargs.get("return_dict", True) else encoded["input_ids"]
tokenizer.apply_chat_template = apply_exact
tokenizer._hessian_exact_template_configured = True
# Parser limits cover a single native generation (currently at most 768 IDs).
# They deliberately do not borrow or modify the experiment's SIGALRM timer.
_QWEN_PARSE_MAX_IDS = 8192
_QWEN_PARSE_MAX_CHARS = 131072
_QWEN_PARSE_SECONDS = 1.0
_QWEN_SCHEMA_SHA256 = "7d16229c9229da54c0364270a739ef83e951d523db2f8643c8197d41ab182c97"
def _parse_qwen3_bounded(tokenizer: Any, ids: list[int], *, schema=None) -> dict[str, Any]:
"""Parse the pinned schema with bounded regex work and exact sampled history.
One lazy body can already span every complete tool block. Removing the
redundant outer repetition preserves its captured substring while avoiding
exponentially many partitions when a later block is unfinished. The native
iterator still parses individual blocks, including its malformed-JSON
failures. An unsupported schema or exhausted limit fails
closed. TimeoutError is never caught: the existing experiment wall guard
must propagate unchanged.
"""
import time
deadline = time.monotonic() + _QWEN_PARSE_SECONDS
import regex
from transformers.utils.chat_parsing_utils import recursive_parse
from trl.chat_template_utils import qwen3_schema
def remaining() -> float:
value = deadline - time.monotonic()
if value <= 0:
raise TimeoutError("Bounded Qwen response parsing exceeded its deadline")
return value
if (not isinstance(ids, list) or len(ids) > _QWEN_PARSE_MAX_IDS
or any(type(token) is not int for token in ids)):
raise ValueError("Qwen response token shape or size is unsupported")
native = qwen3_schema
digest = hashlib.sha256(json.dumps(native, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
if digest != _QWEN_SCHEMA_SHA256 or (schema is not None and schema != native):
raise ValueError("Bounded Qwen parser requires the exact pinned response schema")
text = tokenizer.decode(ids)
if len(text) > _QWEN_PARSE_MAX_CHARS:
raise ValueError("Qwen response character limit exceeded")
# The outer capture needs one span; its body already admits multiple blocks.
old = r"(?:.+?\s*)+"
new = r".+?\s*"
if native["x-regex"].count(old) != 1:
raise ValueError("Pinned Qwen tool-block pattern changed")
# The two engines differ on four C0 whitespace controls. Spell out Python's
# Unicode whitespace set to preserve the pinned stdlib schema semantics.
whitespace = r"[\t\n\r\f\v\x1c-\x1f \x85\xa0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000]"
def native_whitespace(pattern: str) -> str:
return (pattern.replace(r"[\s]", whitespace).replace(r"\s", whitespace)
.replace(r"\S", "[^" + whitespace[1:]))
pattern = native_whitespace(native["x-regex"].replace(old, new))
match = regex.search(pattern, text, flags=regex.DOTALL, timeout=remaining())
if match is None:
raise ValueError("Qwen response did not match the pinned schema")
groups = {key: value for key, value in match.groupdict().items() if value is not None}
plain_schema = copy.deepcopy(native)
del plain_schema["x-regex"]
tool_schema = plain_schema["properties"]["tool_calls"]
tool_pattern = native_whitespace(tool_schema.pop("x-regex-iterator"))
if "tool_calls" in groups:
calls = []
for item in regex.finditer(tool_pattern, groups["tool_calls"], flags=regex.DOTALL, timeout=remaining()):
remaining()
if item.end() <= item.start() or len(item.groups()) != 1:
raise ValueError("Pinned Qwen tool iterator shape changed")
calls.append(item.group(1))
groups["tool_calls"] = calls
remaining()
# This exact pinned remainder contains JSON/JMESPath/type validation only;
# arguments' additionalProperties={} passes values through without recursion.
parsed = recursive_parse(groups, plain_schema)
remaining()
raw = tokenizer.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)
if len(raw) > _QWEN_PARSE_MAX_CHARS:
raise ValueError("Qwen raw response character limit exceeded")
remaining()
parsed["raw_response"] = raw
parsed["raw_token_ids"] = list(ids)
return parsed
def configure_policy_tokenizer(tokenizer: Any) -> dict[str, Any]:
"""Shared training/evaluation Qwen parser and raw assistant history formatting.
Keep parsed messages for tool execution while tokenized histories use actual
sampled assistant IDs. Training independently proves their conditioning
contexts and loss masks before allowing optimization.
"""
from trl.chat_template_utils import (qwen3_chat_template, qwen3_schema,
qwen3_training_chat_template)
template = _raw_preserving_template(qwen3_training_chat_template)
if tokenizer.chat_template not in {qwen3_chat_template, qwen3_training_chat_template, template}:
raise ValueError("Agent pilot supports only the pinned native Qwen3 chat template")
tokenizer.response_schema = qwen3_schema
if not getattr(tokenizer, "_hessian_raw_parser_configured", False):
def parse_with_raw(ids: list[int], **kwargs: Any) -> dict[str, Any]:
return _parse_qwen3_bounded(tokenizer, ids, **kwargs)
tokenizer.parse_response = parse_with_raw
tokenizer._hessian_raw_parser_configured = True
tokenizer.chat_template = template
_install_token_preserving_chat_template(tokenizer)
tokenizer.padding_side = "left"
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
return {"format_version": FORMAT_VERSION, "chat_template": template,
"chat_template_sha256": hashlib.sha256(template.encode()).hexdigest(),
"helper_sha256": sha256(Path(__file__)), "response_schema": qwen3_schema,
"response_parser": {"version": "qwen3-bounded-single-span-v1",
"schema_sha256": _QWEN_SCHEMA_SHA256,
"max_ids": _QWEN_PARSE_MAX_IDS,
"max_characters": _QWEN_PARSE_MAX_CHARS,
"seconds": _QWEN_PARSE_SECONDS}}
def audit_policy_tokens(
prompt_ids: list[list[int]], completion_ids: list[list[int]],
tool_masks: list[list[int]] | None, sampled_spans: list[dict[str, Any]],
) -> list[dict[str, Any]]:
"""Prove each loss token was sampled with exactly the recorded token context.
Decoding then re-encoding is deliberately insufficient. Native retained
token streams must match IDs returned directly by actual model.generate.
External observations/framing must be exactly the complement of those spans.
"""
if tool_masks is None or not (len(prompt_ids) == len(completion_ids) == len(tool_masks)):
raise ValueError("Native tool masks are required for every agent trajectory")
audits = []
for prompt, completion, mask in zip(prompt_ids, completion_ids, tool_masks, strict=True):
if len(completion) != len(mask) or any(type(x) is not int or x not in (0, 1) for x in mask):
raise ValueError("Tool loss mask must align exactly with completion token IDs")
full = prompt + completion
proven = [False] * len(completion)
matched = []
for span in sampled_spans:
context, generated = span["prompt_ids"], span["completion_ids"]
offset = len(context) - len(prompt)
if offset < 0 or offset >= len(completion) or full[:len(context)] != context:
continue
retained = min(len(generated), len(completion) - offset)
if retained < 1 or completion[offset:offset + retained] != generated[:retained]:
continue
proven[offset:offset + retained] = [True] * retained
matched.append({"offset": offset, "retained_tokens": retained,
"prompt_ids": context, "sampled_completion_ids": generated,
"call_index": span["call_index"]})
if any(bool(loss) != proof for loss, proof in zip(mask, proven, strict=True)):
raise ValueError("Native tool stream diverged from actually sampled assistant IDs/context or tool mask")
audits.append({"prompt_ids": prompt, "completion_ids": completion, "tool_mask": mask,
"assistant_tokens": sum(mask), "external_tokens": len(mask) - sum(mask),
"sampled_spans": matched, "sampled_token_equality_verified": True})
return audits
def audit_task_group(tasks: list[dict], environments: list[Any], generations: int) -> list[dict[str, Any]]:
"""Bind each rollout to trusted reset metadata and one common GRPO task."""
if not isinstance(tasks, list) or len(tasks) != generations or len(environments) != generations:
raise ValueError("Task/environment batches must form exactly one complete GRPO group")
metadata = []
for task, environment in zip(tasks, environments, strict=True):
if not isinstance(task, dict) or task != getattr(environment, "_task", None):
raise ValueError("Reward task metadata differs from the environment's actual reset task")
identity = _text(task.get("task_id"), "task.task_id")
episode_seed = task.get("episode_seed")
if type(episode_seed) is not int or episode_seed < 0 or task.get("split") != "train":
raise ValueError("Task metadata requires a train split and nonnegative integer episode seed")
metadata.append({"task_id": identity, "episode_seed": episode_seed, "split": "train"})
if len({(row["task_id"], row["episode_seed"]) for row in metadata}) != 1:
raise ValueError("All rollouts in a GRPO group must use the same task identity and episode seed")
return metadata
def _frozen_file(base: Path, value: Any, label: str) -> tuple[Path, str]:
if not isinstance(value, dict):
raise ValueError(f"{label} must be a frozen path/sha256 object")
_keys(value, {"path", "sha256"}, label)
path = (base / _text(value.get("path"), f"{label}.path")).resolve()
expected = _text(value.get("sha256"), f"{label}.sha256")
actual = checkpoint_sha256(path) if label == "model" else sha256(path)
if expected != actual:
raise ValueError(f"{label} SHA256 mismatch")
return path, actual
def plan(config_path: str | Path) -> dict[str, Any]:
path = Path(config_path).resolve()
config, digest = read_json_snapshot(path)
_keys(config, {"schema_version", "model", "dataset", "output_dir", "seed",
"chat_template_kwargs", "training", "lora"}, "agent RL recipe")
if type(config.get("schema_version")) is not int or config["schema_version"] != 1:
raise ValueError("schema_version must be 1")
if int(os.environ.get("WORLD_SIZE", "1")) != 1:
raise ValueError("The bounded native agent pilot supports one CUDA process only")
model, model_hash = _frozen_file(path.parent, config.get("model"), "model")
if not model.is_dir() or (model / "adapter_config.json").exists():
raise ValueError("model must be a standalone local full checkpoint")
dataset, data_hash = _frozen_file(path.parent, config.get("dataset"), "dataset")
rows, snapshot_hash = read_jsonl_snapshot(dataset)
if snapshot_hash != data_hash:
raise ValueError("Dataset changed while validating")
ids = set()
for row in rows:
identity = _text(row.get("id"), "task id")
_text(row.get("prompt"), "task prompt")
if identity in ids:
raise ValueError("Duplicate agent training task id")
ids.add(identity)
if not isinstance(row.get("task"), dict) or row["task"].get("split") != "train":
raise ValueError("Agent training accepts only explicit train-split task metadata")
settings = config.get("training", {})
if not isinstance(settings, dict):
raise ValueError("training must be an object")
_keys(settings, set(DEFAULT_TRAINING), "agent training")
settings = DEFAULT_TRAINING | settings
integers = {"max_steps", "num_generations", "per_device_train_batch_size",
"gradient_accumulation_steps", "max_completion_length", "max_tool_calling_iterations", "top_k"}
for key, value in settings.items():
_positive(value, key, integer=key in integers, zero=key in {"beta", "top_k", "min_p"})
if settings["num_generations"] < 2 or settings["per_device_train_batch_size"] != settings["num_generations"]:
raise ValueError("Use exactly one grouped prompt: train batch must equal num_generations >= 2")
if settings["gradient_accumulation_steps"] != 1:
raise ValueError("Agent pilot fixes accumulation=1 so each update uses freshly sampled grouped rollouts")
if settings["top_p"] > 1 or settings["min_p"] > 1:
raise ValueError("top_p and min_p must be <= 1")
if settings["max_steps"] > 256 or settings["num_generations"] > 16 or settings["max_tool_calling_iterations"] > 8:
raise ValueError("Agent pilot exceeds bounded step/rollout/tool-round limits")
lora = config.get("lora", DEFAULT_LORA)
if not isinstance(lora, dict):
raise ValueError("Agent pilot requires a LoRA configuration")
_keys(lora, set(DEFAULT_LORA), "LoRA")
lora = DEFAULT_LORA | lora
_positive(lora["r"], "lora.r", integer=True)
_positive(lora["alpha"], "lora.alpha", integer=True)
_positive(lora["dropout"], "lora.dropout", zero=True)
if lora["dropout"] != 0 or lora["target_modules"] != "all-linear":
raise ValueError("On-policy pilot fixes LoRA dropout=0 and target_modules=all-linear")
mode = chat_template_kwargs(config.get("chat_template_kwargs", {"enable_thinking": True}))
output = (path.parent / _text(config.get("output_dir"), "output_dir")).resolve()
if output.exists() or output.is_relative_to(model) or model.is_relative_to(output) or output in dataset.parents:
raise ValueError("Agent output must be fresh and separate from model/data")
seed = config.get("seed", 42)
_positive(seed, "seed", integer=True, zero=True)
return {"schema_version": 1, "status": "planned", "config_path": str(path),
"config_sha256": digest, "model": {"path": str(model), "sha256": model_hash},
"dataset": {"path": str(dataset), "sha256": data_hash}, "rows": len(rows),
"output_dir": str(output), "seed": seed, "chat_template_kwargs": mode,
"training": settings, "lora": lora, "method": "native_tool_environment_grpo",
"fixed": {"full_determinism": True, "bf16": True, "gradient_checkpointing": True, "num_iterations": 1,
"steps_per_generation": 1, "loss_type": "grpo", "use_vllm": False,
"report_to": "none", "push_to_hub": False},
"reward_protocol": dict(REWARD_PROTOCOL), "helper_sha256": sha256(Path(__file__))}
def _audited_trainer_class(base_class: type) -> type:
class AuditedNativeGRPOTrainer(base_class):
def _generate_single_turn(self, prompts: list) -> tuple:
result = super()._generate_single_turn(prompts)
call_index = len(self._sampled_calls)
self._sampled_calls.extend(
{"prompt_ids": list(p), "completion_ids": list(c), "call_index": call_index}
for p, c in zip(result[0], result[1], strict=True)
)
return result
def _generate(self, prompts: list) -> tuple:
self._sampled_calls = []
result = super()._generate(prompts)
try:
self._last_policy_audit = audit_policy_tokens(result[0], result[1], result[2], self._sampled_calls)
except ValueError as error:
failure_path = Path(self.args.output_dir) / "rollout-audit-failure.json"
write_manifest(failure_path, {"error": str(error), "prompt_ids": result[0],
"completion_ids": result[1], "tool_masks": result[2],
"sampled_spans": self._sampled_calls})
raise
return result
return AuditedNativeGRPOTrainer
def run(config_path: str | Path) -> dict[str, Any]:
report = plan(config_path)
rows, digest = read_jsonl_snapshot(report["dataset"]["path"])
if digest != report["dataset"]["sha256"] or sha256(config_path) != report["config_sha256"]:
raise ValueError("Recipe or tasks changed after validation")
installed = require_versions()
if importlib.metadata.version("jmespath") != "1.0.1":
raise RuntimeError("Native tools require the pinned jmespath==1.0.1")
import torch
from datasets import Dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
from trl import GRPOConfig, GRPOTrainer
from . import agent_rl_tasks
if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported() or torch.cuda.device_count() != 1:
raise RuntimeError("Use a single visible CUDA GPU allocation with BF16 support")
output = Path(report["output_dir"])
output.mkdir(parents=True, exist_ok=False)
evidence_path = output / "rollout-evidence.jsonl"
manifest_path = output / "training-manifest.json"
report.update(status="running", dependencies=installed | {"jmespath": "1.0.1"},
environment_sha256=sha256(Path(agent_rl_tasks.__file__)),
upstream_trainer_sha256=sha256(Path(inspect.getfile(GRPOTrainer))),
started_at=datetime.now(timezone.utc).isoformat(), gpu=torch.cuda.get_device_name(),
evidence=str(evidence_path), reference_policy="merged SFT base with new RL LoRA disabled")
write_manifest(manifest_path, report)
try:
set_seed(report["seed"])
tokenizer = AutoTokenizer.from_pretrained(report["model"]["path"], local_files_only=True, trust_remote_code=False)
report["policy_format"] = configure_policy_tokenizer(tokenizer)
tokenizer.save_pretrained(str(output / "policy-tokenizer"))
model = AutoModelForCausalLM.from_pretrained(
report["model"]["path"], local_files_only=True, trust_remote_code=False,
use_safetensors=True, dtype=torch.bfloat16,
)
if checkpoint_sha256(report["model"]["path"]) != report["model"]["sha256"]:
raise ValueError("Base checkpoint changed while loading")
model.config.use_cache = False
model.generation_config.eos_token_id = tokenizer.eos_token_id
settings = dict(report["training"])
reward_groups: list[dict[str, Any]] = []
Trainer = _audited_trainer_class(GRPOTrainer)
trainer = None
def environment_reward(completions: list, environments: list, **kwargs: Any) -> list[float]:
if trainer is None or len(completions) != len(environments):
raise ValueError("Environment and rollout batches must align")
metadata = audit_task_group(kwargs.get("task"), environments, settings["num_generations"])
policy_version = {"global_step": trainer.state.global_step,
"base_model_sha256": report["model"]["sha256"]}
scores = [agent_rl_tasks.score_episode(env, completion)
for env, completion in zip(environments, completions, strict=True)]
components = [milestone_training_reward(score) for score in scores]
rewards = [component["training_reward"] for component in components]
if any(not math.isfinite(r) or r < 0 or r > 1 for r in rewards):
raise ValueError("Environment reward must be finite in [0,1]")
group = {"optimizer_step": trainer.state.global_step, "rewards": rewards,
"reward_std": statistics.pstdev(rewards),
"successes": sum(bool(score["success"]) for score in scores),
"strict_successes": sum(bool(score["success"]) for score in scores),
"strict_rewards": [float(score["reward"]) for score in scores],
"task_ids": [row["task_id"] for row in metadata],
"episode_seeds": [row["episode_seed"] for row in metadata],
"distinct_task_count": 1, "policy_version": policy_version}
reward_groups.append(group)
audits = trainer._last_policy_audit
with evidence_path.open("a", encoding="utf-8") as stream:
for index, (score, completion, audit) in enumerate(zip(scores, completions, audits, strict=True)):
stream.write(json.dumps({"optimizer_step": trainer.state.global_step,
"rollout_index": index, "score": score,
"training_reward": rewards[index],
"reward_components": components[index],
**metadata[index], "policy_version": policy_version,
"completion_messages": completion, "policy_audit": audit},
ensure_ascii=False, allow_nan=False) + "\n")
return rewards
arguments = GRPOConfig(
output_dir=str(output), **settings, num_iterations=1, steps_per_generation=1,
loss_type="grpo", full_determinism=True, bf16=True, fp16=False, gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": False},
chat_template_kwargs=report["chat_template_kwargs"], use_vllm=False,
disable_dropout=True, report_to="none", push_to_hub=False,
seed=report["seed"], data_seed=report["seed"], optim="adamw_torch",
save_strategy="no", logging_steps=1, dataloader_pin_memory=False,
mask_truncated_completions=False, log_completions=False,
)
lora = report["lora"]
trainer = Trainer(
model=model, processing_class=tokenizer, args=arguments,
train_dataset=Dataset.from_list([
{"id": row["id"], "prompt": [{"role": "user", "content": row["prompt"]}], "task": row["task"]}
for row in rows
]), reward_funcs=environment_reward, environment_factory=agent_rl_tasks.GermanAgentEnv,
peft_config=LoraConfig(r=lora["r"], lora_alpha=lora["alpha"], lora_dropout=0.0,
target_modules="all-linear", task_type="CAUSAL_LM"),
)
if trainer.ref_model is not None:
raise ValueError("Expected LoRA-disabled merged SFT reference, not an untracked reference model")
before = {name: tensor.detach().cpu().clone() for name, tensor in trainer.model.named_parameters()
if tensor.requires_grad}
if not before:
raise ValueError("LoRA created no trainable parameters")
trained = trainer.train()
changed, max_delta = [], 0.0
for name, tensor in trainer.model.named_parameters():
if name in before:
current = tensor.detach().cpu()
if not torch.isfinite(current).all():
raise ValueError("Non-finite trained adapter parameter")
delta = float((current.float() - before[name].float()).abs().max())
if delta > 0:
changed.append(name)
max_delta = max(max_delta, delta)
trainer.save_model(str(output / "checkpoint"))
tokenizer.save_pretrained(str(output / "checkpoint"))
report.update(status="completed", checkpoint=str(output / "checkpoint"),
checkpoint_sha256=checkpoint_sha256(output / "checkpoint"), metrics=trained.metrics,
train_log_history=trainer.state.log_history, reward_groups=reward_groups,
task_coverage={"training_pool_rows": len(rows),
"consumed_unique_task_ids": sorted({identity for group in reward_groups for identity in group["task_ids"]}),
"total_rollouts": sum(len(group["task_ids"]) for group in reward_groups),
"optimizer_groups": len(reward_groups)},
update_evidence={"changed_trainable_tensors": len(changed), "max_abs_parameter_delta": max_delta,
"mixed_reward_groups": sum(group["reward_std"] > 0 for group in reward_groups),
"nonzero_gradient_steps": sum(float(log.get("grad_norm", 0)) > 0 for log in trainer.state.log_history)},
completed_at=datetime.now(timezone.utc).isoformat())
if evidence_path.exists():
report["evidence_sha256"] = sha256(evidence_path)
write_manifest(manifest_path, report)
return report
except BaseException as error:
report.update(status="failed", error={"type": type(error).__name__, "message": str(error)})
write_manifest(manifest_path, report)
raise
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("config", type=Path)
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
try:
print(json.dumps(plan(args.config) if args.dry_run else run(args.config), indent=2))
except (ValueError, OSError, RuntimeError, ImportError) as error:
parser.exit(2, f"error: {error}\n")
if __name__ == "__main__":
main()