skillrl 1.0.0__py3-none-any.whl

skillrl/__init__.py

@@ -0,0 +1,49 @@
+"""skillrl — A TRL-like training library for end-to-end skill optimization.
+
+This package implements the core algorithm of Microsoft SkillOpt
+(https://microsoft.github.io/SkillOpt/) as a clean, modular, TRL-style
+training library.  The trainable state is a natural-language *skill
+document*; both the optimizer and the target LLM stay frozen.
+
+Quick start
+-----------
+>>> from skillrl import SkillOptConfig, SkillOptTrainer
+>>> from skillrl.envs.qa import SimpleQAEnv
+>>> from skillrl.llm.openai_client import OpenAIChatClient
+>>>
+>>> cfg = SkillOptConfig(num_epochs=2, batch_size=8, edit_budget=4)
+>>> env = SimpleQAEnv(train_items=[...], val_items=[...], test_items=[...])
+>>> client = OpenAIChatClient(model="gpt-4o-mini")
+>>> trainer = SkillOptTrainer(
+...     config=cfg, env=env,
+...     optimizer_client=client, target_client=client,
+...     initial_skill="You are a helpful assistant.",
+... )
+>>> summary = trainer.train()
+>>> print(summary["test_hard"])
+"""
+
+from skillrl.config import SkillOptConfig
+from skillrl.types import (
+    Edit,
+    Patch,
+    RolloutResult,
+    GateResult,
+    GateAction,
+    RawPatch,
+)
+from skillrl.trainer import SkillOptTrainer
+
+__version__ = "1.0.0"
+
+__all__ = [
+    "__version__",
+    "SkillOptConfig",
+    "SkillOptTrainer",
+    "Edit",
+    "Patch",
+    "RolloutResult",
+    "GateResult",
+    "GateAction",
+    "RawPatch",
+]

skillrl/config.py

@@ -0,0 +1,148 @@
+"""Training configuration for skillrl — a TRL-like dataclass.
+
+Mirrors the paper-default protocol of SkillOpt (Yang et al., 2026) while
+exposing only the knobs needed for the 1.0 core algorithm.  Advanced
+features such as slow update, meta skill, autonomous LR, accumulation,
+codex/claude-code execution backends and full-rewrite update modes are
+intentionally omitted from 1.0 — they are future extension points.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, asdict
+from typing import Literal, Any
+
+
+GateMetric = Literal["hard", "soft", "mixed"]
+LRSchedulerMode = Literal["constant", "linear", "cosine"]
+
+
+@dataclass
+class SkillOptConfig:
+    """Training arguments for :class:`~skillrl.trainer.SkillOptTrainer`.
+
+    The defaults reproduce the SkillOpt paper protocol: 4 epochs,
+    rollout batch 40, reflection minibatch 8, textual learning rate 4
+    with cosine decay, strict hard-validation gating.
+
+    Parameters
+    ----------
+    num_epochs:
+        Number of training epochs.
+    batch_size:
+        Number of training items rolled out per optimization step.
+    minibatch_size:
+        Number of trajectories grouped per analyst (Reflect) call.
+    merge_batch_size:
+        Number of patches merged together at each level of the
+        hierarchical Aggregate stage.
+    edit_budget:
+        Initial / maximum textual learning rate — the cap on edits
+        applied per optimization step.  Acts as a soft trust region.
+    min_edit_budget:
+        Lower bound for decay schedules.
+    lr_scheduler:
+        Edit-budget schedule: ``constant`` | ``linear`` | ``cosine``.
+    gate_metric:
+        Validation-gate metric: ``hard`` (exact match) | ``soft``
+        (continuous reward) | ``mixed``.
+    gate_mixed_weight:
+        Weight ``w`` of the soft component when ``gate_metric='mixed'``.
+        Final score = ``(1 - w) * hard + w * soft``.
+    failure_only:
+        If True, only failed trajectories drive analyst patches.
+    workers:
+        Thread workers for parallel rollout / parallel analyst calls.
+    seed:
+        Random seed.
+    out_root:
+        Output directory.  History, per-step artifacts, candidate
+        skills and ``best_skill.md`` are written here.  Auto-resume
+        reads from this directory if it already contains a run.
+    selection_split:
+        Name of the split used for the validation gate.  Defaults to
+        ``"val"`` (the SkillOpt paper uses ``valid_seen``).
+    test_split:
+        Name of the held-out test split for the final report.
+    eval_test:
+        Whether to run the final test evaluation at the end of
+        training.
+    save_every_step:
+        If True, dump per-step ``skill_vXXXX.md`` snapshots to
+        ``<out_root>/skills/`` for inspection.
+    optimizer_max_completion_tokens:
+        Token budget for analyst / merge / ranking calls.
+    verbose:
+        Print detailed per-stage progress.
+    extras:
+        Free-form extra config — passed through to environments and
+        custom prompt templates.
+    """
+
+    # ── Training schedule ────────────────────────────────────────────
+    num_epochs: int = 4
+    batch_size: int = 40
+    minibatch_size: int = 8
+    merge_batch_size: int = 8
+
+    # ── Optimizer (textual LR) ───────────────────────────────────────
+    edit_budget: int = 4
+    min_edit_budget: int = 2
+    lr_scheduler: LRSchedulerMode = "cosine"
+
+    # ── Validation gate ──────────────────────────────────────────────
+    gate_metric: GateMetric = "hard"
+    gate_mixed_weight: float = 0.5
+
+    # ── Reflect stage ────────────────────────────────────────────────
+    failure_only: bool = False
+
+    # ── Concurrency / runtime ────────────────────────────────────────
+    workers: int = 8
+    seed: int = 42
+    out_root: str = "outputs/skillrl_run"
+
+    # ── Evaluation splits ────────────────────────────────────────────
+    selection_split: str = "val"
+    test_split: str = "test"
+    eval_test: bool = True
+
+    # ── Persistence / observability ──────────────────────────────────
+    save_every_step: bool = True
+    optimizer_max_completion_tokens: int = 4096
+    verbose: bool = True
+
+    # ── Free-form extras (forwarded to env/prompts) ──────────────────
+    extras: dict[str, Any] = field(default_factory=dict)
+
+    # ── Helpers ──────────────────────────────────────────────────────
+
+    def __post_init__(self) -> None:
+        if self.batch_size <= 0:
+            raise ValueError(f"batch_size must be > 0, got {self.batch_size}")
+        if self.minibatch_size <= 0:
+            raise ValueError(f"minibatch_size must be > 0, got {self.minibatch_size}")
+        if self.merge_batch_size < 2:
+            raise ValueError(f"merge_batch_size must be >= 2, got {self.merge_batch_size}")
+        if self.edit_budget <= 0:
+            raise ValueError(f"edit_budget must be > 0, got {self.edit_budget}")
+        if self.min_edit_budget <= 0 or self.min_edit_budget > self.edit_budget:
+            raise ValueError(
+                f"min_edit_budget must be in (0, edit_budget], "
+                f"got {self.min_edit_budget} (edit_budget={self.edit_budget})"
+            )
+        if self.lr_scheduler not in ("constant", "linear", "cosine"):
+            raise ValueError(
+                f"lr_scheduler must be one of constant/linear/cosine, "
+                f"got {self.lr_scheduler!r}"
+            )
+        if self.gate_metric not in ("hard", "soft", "mixed"):
+            raise ValueError(
+                f"gate_metric must be one of hard/soft/mixed, got {self.gate_metric!r}"
+            )
+        if not 0.0 <= self.gate_mixed_weight <= 1.0:
+            raise ValueError(
+                f"gate_mixed_weight must be in [0, 1], got {self.gate_mixed_weight}"
+            )
+
+    def to_dict(self) -> dict:
+        return asdict(self)

skillrl/core/__init__.py

@@ -0,0 +1,36 @@
+"""skillrl core algorithms — pure-functional building blocks.
+
+This subpackage hosts the *non-LLM* primitives used by the trainer:
+
+* :mod:`skillrl.core.editor`     — apply edits / patches to a skill doc.
+* :mod:`skillrl.core.scheduler`  — edit-budget (textual LR) schedulers.
+* :mod:`skillrl.core.gate`       — validation gating (hard/soft/mixed).
+* :mod:`skillrl.core.utils`      — JSON extraction, scoring, hashing.
+"""
+
+from skillrl.core.editor import apply_edit, apply_patch, apply_patch_with_report
+from skillrl.core.scheduler import (
+    LRScheduler,
+    ConstantScheduler,
+    LinearScheduler,
+    CosineScheduler,
+    build_scheduler,
+)
+from skillrl.core.gate import evaluate_gate, select_gate_score
+from skillrl.core.utils import compute_score, extract_json, skill_hash
+
+__all__ = [
+    "apply_edit",
+    "apply_patch",
+    "apply_patch_with_report",
+    "LRScheduler",
+    "ConstantScheduler",
+    "LinearScheduler",
+    "CosineScheduler",
+    "build_scheduler",
+    "evaluate_gate",
+    "select_gate_score",
+    "compute_score",
+    "extract_json",
+    "skill_hash",
+]

skillrl/core/editor.py

@@ -0,0 +1,110 @@
+"""Skill-document edit application — the Update stage (⑤).
+
+Analogous to ``optimizer.step()`` in neural network training: a ranked
+:class:`~skillrl.types.Patch` is applied sequentially onto the current
+skill document, producing a candidate skill that is then validated by
+the gate.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from skillrl.types import Edit, Patch
+
+
+def _edit_fields(edit: Edit | dict) -> tuple[str, str, str]:
+    """Extract ``(op, content, target)`` from an ``Edit`` or plain dict."""
+    if isinstance(edit, Edit):
+        return edit.op, (edit.content or "").strip(), edit.target or ""
+    return (
+        str(edit.get("op", "") or ""),
+        str(edit.get("content", "") or "").strip(),
+        str(edit.get("target", "") or ""),
+    )
+
+
+def _apply_edit_with_report(skill: str, edit: Edit | dict) -> tuple[str, dict]:
+    op, content, target = _edit_fields(edit)
+    report: dict[str, Any] = {
+        "op": op,
+        "target": target[:200],
+        "content_preview": content[:200],
+        "status": "unknown",
+    }
+
+    if op == "append":
+        report["status"] = "applied_append"
+        return skill.rstrip() + "\n\n" + content + "\n", report
+
+    if op == "insert_after":
+        if not target or target not in skill:
+            report["status"] = "applied_insert_after_fallback_append"
+            return skill.rstrip() + "\n\n" + content + "\n", report
+        idx = skill.index(target) + len(target)
+        newline = skill.find("\n", idx)
+        insert_at = newline + 1 if newline != -1 else len(skill)
+        report["status"] = "applied_insert_after"
+        return skill[:insert_at] + "\n" + content + "\n" + skill[insert_at:], report
+
+    if op == "replace":
+        if not target:
+            report["status"] = "skipped_replace_missing_target"
+            return skill, report
+        if target not in skill:
+            report["status"] = "skipped_replace_target_not_found"
+            return skill, report
+        report["status"] = "applied_replace"
+        return skill.replace(target, content, 1), report
+
+    if op == "delete":
+        if not target:
+            report["status"] = "skipped_delete_missing_target"
+            return skill, report
+        if target not in skill:
+            report["status"] = "skipped_delete_target_not_found"
+            return skill, report
+        report["status"] = "applied_delete"
+        return skill.replace(target, "", 1), report
+
+    report["status"] = "skipped_unknown_op"
+    return skill, report
+
+
+def apply_edit(skill: str, edit: Edit | dict) -> str:
+    """Apply a single edit operation to the skill document."""
+    updated, _ = _apply_edit_with_report(skill, edit)
+    return updated
+
+
+def apply_patch_with_report(
+    skill: str,
+    patch: Patch | dict,
+) -> tuple[str, list[dict]]:
+    """Apply a patch and return a per-edit report for observability."""
+    if isinstance(patch, Patch):
+        edits: list[Edit | dict] = list(patch.edits)
+    else:
+        edits = list(patch.get("edits", []) or [])
+
+    reports: list[dict] = []
+    for idx, edit in enumerate(edits, 1):
+        try:
+            skill, report = _apply_edit_with_report(skill, edit)
+            report["index"] = idx
+        except Exception as exc:  # noqa: BLE001
+            report = {
+                "index": idx,
+                "op": "",
+                "target": "",
+                "content_preview": "",
+                "status": "error",
+                "error": str(exc),
+            }
+        reports.append(report)
+    return skill, reports
+
+
+def apply_patch(skill: str, patch: Patch | dict) -> str:
+    """Apply a patch (list of edits) to the skill document sequentially."""
+    updated, _ = apply_patch_with_report(skill, patch)
+    return updated

skillrl/core/gate.py

@@ -0,0 +1,94 @@
+"""Validation gating — Stage ⑥ of the skillrl pipeline.
+
+A candidate skill is *accepted* only when its score on the held-out
+selection split *strictly improves* upon the running ``current_score``.
+The metric is configurable:
+
+* ``hard``   : exact-match accuracy (paper default).
+* ``soft``   : continuous per-item reward.
+* ``mixed``  : ``(1 - w) * hard + w * soft``.
+
+The gate also tracks the global ``best_skill`` separately, so the
+trainer can deploy the all-time-best skill at the end of training.
+"""
+from __future__ import annotations
+
+from typing import Literal
+
+from skillrl.types import GateResult
+
+
+GateMetric = Literal["hard", "soft", "mixed"]
+
+
+def select_gate_score(
+    hard: float,
+    soft: float,
+    metric: GateMetric = "hard",
+    mixed_weight: float = 0.5,
+) -> float:
+    """Combine ``(hard, soft)`` into a single gate score per ``metric``."""
+    if metric == "hard":
+        return float(hard)
+    if metric == "soft":
+        return float(soft)
+    if metric == "mixed":
+        w = float(mixed_weight)
+        if not 0.0 <= w <= 1.0:
+            raise ValueError(f"mixed_weight must be in [0, 1], got {w}")
+        return (1.0 - w) * float(hard) + w * float(soft)
+    raise ValueError(f"Unknown gate metric: {metric!r}")
+
+
+def evaluate_gate(
+    *,
+    candidate_skill: str,
+    cand_hard: float,
+    current_skill: str,
+    current_score: float,
+    best_skill: str,
+    best_score: float,
+    best_step: int,
+    global_step: int,
+    cand_soft: float = 0.0,
+    metric: GateMetric = "hard",
+    mixed_weight: float = 0.5,
+) -> GateResult:
+    """Decide accept / reject and update ``current`` / ``best`` skill state.
+
+    The decision rule is:
+
+    * If ``cand_score > best_score`` → ``accept_new_best``: replace
+      both ``current`` and ``best``.
+    * Else if ``cand_score > current_score`` → ``accept``: replace
+      ``current`` only.
+    * Else → ``reject``: keep both unchanged.
+    """
+    cand_score = select_gate_score(cand_hard, cand_soft, metric, mixed_weight)
+
+    if cand_score > best_score:
+        return GateResult(
+            action="accept_new_best",
+            current_skill=candidate_skill,
+            current_score=cand_score,
+            best_skill=candidate_skill,
+            best_score=cand_score,
+            best_step=global_step,
+        )
+    if cand_score > current_score:
+        return GateResult(
+            action="accept",
+            current_skill=candidate_skill,
+            current_score=cand_score,
+            best_skill=best_skill,
+            best_score=best_score,
+            best_step=best_step,
+        )
+    return GateResult(
+        action="reject",
+        current_skill=current_skill,
+        current_score=current_score,
+        best_skill=best_skill,
+        best_score=best_score,
+        best_step=best_step,
+    )

skillrl/core/scheduler.py

@@ -0,0 +1,88 @@
+"""Edit-budget (textual learning rate) schedulers.
+
+In skillrl the *learning rate* is the maximum number of edits applied
+to the skill document at each optimization step.  Exposing the same
+PyTorch-style API (``step()``, ``state_dict()`` / ``load_state_dict``)
+keeps the abstraction familiar.
+"""
+from __future__ import annotations
+
+import math
+from abc import ABC, abstractmethod
+
+
+class LRScheduler(ABC):
+    """Base class for edit-budget schedulers."""
+
+    def __init__(self, max_lr: int, min_lr: int, total_steps: int) -> None:
+        self.max_lr = int(max_lr)
+        self.min_lr = int(min_lr)
+        self.total_steps = max(int(total_steps), 1)
+        self._current_step = 0
+
+    @abstractmethod
+    def _compute_lr(self, step: int) -> int:
+        """Return the edit budget for a 1-indexed step."""
+
+    def step(self) -> int:
+        self._current_step += 1
+        return self._compute_lr(self._current_step)
+
+    def get_lr(self, step: int) -> int:
+        return self._compute_lr(step)
+
+    def state_dict(self) -> dict:
+        return {"current_step": self._current_step}
+
+    def load_state_dict(self, state: dict) -> None:
+        self._current_step = int(state.get("current_step", 0) or 0)
+
+
+class ConstantScheduler(LRScheduler):
+    """Fixed budget throughout training."""
+
+    def _compute_lr(self, step: int) -> int:
+        return self.max_lr
+
+
+class LinearScheduler(LRScheduler):
+    """Linear decay from ``max_lr`` to ``min_lr`` over ``total_steps``."""
+
+    def _compute_lr(self, step: int) -> int:
+        if self.total_steps <= 1:
+            return self.max_lr
+        t = min(step, self.total_steps) / self.total_steps
+        lr = self.max_lr + (self.min_lr - self.max_lr) * t
+        return max(self.min_lr, round(lr))
+
+
+class CosineScheduler(LRScheduler):
+    """Cosine annealing from ``max_lr`` to ``min_lr`` over ``total_steps``."""
+
+    def _compute_lr(self, step: int) -> int:
+        if self.total_steps <= 1:
+            return self.max_lr
+        t = min(step, self.total_steps) / self.total_steps
+        lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + math.cos(math.pi * t))
+        return max(self.min_lr, round(lr))
+
+
+_REGISTRY: dict[str, type[LRScheduler]] = {
+    "constant": ConstantScheduler,
+    "linear": LinearScheduler,
+    "cosine": CosineScheduler,
+}
+
+
+def build_scheduler(
+    mode: str = "cosine",
+    max_lr: int = 4,
+    min_lr: int = 2,
+    total_steps: int = 8,
+) -> LRScheduler:
+    """Build an edit-budget scheduler from config parameters."""
+    if mode not in _REGISTRY:
+        raise ValueError(
+            f"Unknown scheduler mode {mode!r}. Available: {list(_REGISTRY)}"
+        )
+    return _REGISTRY[mode](max_lr=max_lr, min_lr=min_lr, total_steps=total_steps)

skillrl/core/utils.py

@@ -0,0 +1,96 @@
+"""Misc utilities — JSON extraction, scoring, hashing."""
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from typing import Any, Iterable
+
+from skillrl.types import RolloutResult
+
+
+# ── Robust JSON extraction from LLM output ──────────────────────────────
+
+
+_FENCE_RE = re.compile(r"```(?:json)?\s*(.+?)```", re.DOTALL | re.IGNORECASE)
+
+
+def extract_json(text: str) -> dict | None:
+    """Best-effort extraction of a JSON object embedded in *text*.
+
+    Tries (in order):
+
+    1. parsing the entire string as JSON,
+    2. parsing the contents of the first fenced code block,
+    3. parsing the substring between the first ``{`` and last ``}``.
+
+    Returns ``None`` on failure — callers must handle this gracefully.
+    """
+    if text is None:
+        return None
+    s = text.strip()
+    if not s:
+        return None
+
+    # 1. raw JSON
+    try:
+        obj = json.loads(s)
+        return obj if isinstance(obj, dict) else None
+    except Exception:
+        pass
+
+    # 2. fenced block
+    m = _FENCE_RE.search(s)
+    if m:
+        try:
+            obj = json.loads(m.group(1).strip())
+            return obj if isinstance(obj, dict) else None
+        except Exception:
+            pass
+
+    # 3. brace slice (last-resort)
+    first = s.find("{")
+    last = s.rfind("}")
+    if first != -1 and last != -1 and last > first:
+        candidate = s[first : last + 1]
+        try:
+            obj = json.loads(candidate)
+            return obj if isinstance(obj, dict) else None
+        except Exception:
+            pass
+
+    return None
+
+
+# ── Aggregate rollout scoring ───────────────────────────────────────────
+
+
+def compute_score(results: Iterable[Any]) -> tuple[float, float]:
+    """Return ``(hard_acc, soft_acc)`` averaged over a rollout batch.
+
+    Accepts ``RolloutResult`` instances or plain dicts (with ``hard``
+    and ``soft`` keys).  Empty input yields ``(0.0, 0.0)``.
+    """
+    items = list(results)
+    if not items:
+        return 0.0, 0.0
+    hard_total = 0.0
+    soft_total = 0.0
+    for r in items:
+        if isinstance(r, RolloutResult):
+            hard_total += float(r.hard)
+            soft_total += float(r.soft)
+        else:
+            hard_total += float(r.get("hard", 0) or 0)
+            soft_total += float(r.get("soft", 0.0) or 0.0)
+    n = len(items)
+    return hard_total / n, soft_total / n
+
+
+# ── Skill content hashing (cache key for selection eval) ────────────────
+
+
+def skill_hash(content: str) -> str:
+    """Deterministic short hash for a skill document (cache key)."""
+    h = hashlib.sha256((content or "").encode("utf-8")).hexdigest()
+    return h[:16]

skillrl/envs/__init__.py

@@ -0,0 +1,19 @@
+"""Environment / dataset abstractions used by the trainer.
+
+A ``SkillEnv`` defines:
+
+* a train / val / test split,
+* how to roll out a single item with a given skill (calling the target
+  LLM via the supplied client),
+* how to score the rollout (``hard``: 0/1 exact-match, ``soft``: [0,1]
+  partial credit).
+
+Bring-your-own-env: subclass :class:`SkillEnv` and pass the instance to
+:class:`~skillrl.SkillOptTrainer`.  See :class:`~skillrl.envs.qa.SimpleQAEnv`
+for a complete reference implementation.
+"""
+
+from skillrl.envs.base import SkillEnv
+from skillrl.envs.qa import SimpleQAEnv
+
+__all__ = ["SkillEnv", "SimpleQAEnv"]