diff --git a/configs/alphabet_sort.toml b/configs/alphabet_sort.toml new file mode 100644 index 0000000000..db01cdbc08 --- /dev/null +++ b/configs/alphabet_sort.toml @@ -0,0 +1,15 @@ +# alphabet-sort-v1 — multi-turn alphabetical sorting, driven by a colocated user simulator. +# +# uv run eval @ configs/alphabet_sort.toml --model +num_tasks = 5 +num_rollouts = 2 +max_turns = 4 + +[taskset] +id = "alphabet-sort-v1" +min_turns = 2 +max_turns = 2 + +[harness] +id = "default" +runtime = { type = "subprocess" } diff --git a/configs/textarena.toml b/configs/textarena.toml index f20f83f1f8..e423bf12dd 100644 --- a/configs/textarena.toml +++ b/configs/textarena.toml @@ -15,5 +15,4 @@ num_tasks = 50 # pool of seeded episodes to draw the eval's tasks from [harness] id = "default" -enable_bash = false runtime = { type = "subprocess" } diff --git a/configs/wordle.toml b/configs/wordle.toml index 4429421ac7..5b49545d04 100644 --- a/configs/wordle.toml +++ b/configs/wordle.toml @@ -11,5 +11,4 @@ num_tasks = 50 # pool of seeded episodes to draw the eval's tasks from [harness] id = "default" -enable_bash = false runtime = { type = "subprocess" } diff --git a/examples/tasksets/alphabet_sort_v1/alphabet_sort_v1/__init__.py b/examples/tasksets/alphabet_sort_v1/alphabet_sort_v1/__init__.py new file mode 100644 index 0000000000..44494b5e9d --- /dev/null +++ b/examples/tasksets/alphabet_sort_v1/alphabet_sort_v1/__init__.py @@ -0,0 +1,186 @@ +"""alphabet-sort-v1 — maintain an alphabetically sorted list of names across turns. + +The v1 port of the `alphabet-sort` environment. Each task is a multi-turn episode: the model +sorts an initial list of names (by first or last name) into `` tags, then +on each follow-up turn re-sorts the cumulative list — tagging the newly added names — into +`` tags. The reward is the per-turn sequence similarity to the +ground truth, power-scaled. + +The follow-up turns are colocated with the agent as a `vf.User` (see `user.py`): the +interception server drives the simulator after every assistant turn and injects the next +follow-up as a user message, so the whole episode is one rollout the harness only ever sees as +a single exchange. The episodes are pre-generated in `load_tasks`; the simulator replays them. +""" + +import difflib +import json +import random +import re +import sys +from typing import Literal + +from datasets import load_dataset + +import verifiers.v1 as vf + +DATASET = "kalomaze/alphabetic-arxiv-authors-it1" +SEED = 1337420 + + +class AlphabetSortConfig(vf.TasksetConfig): + min_turns: int = 1 + """Minimum number of turns (assistant sorts) per episode.""" + max_turns: int = 3 + """Maximum number of turns per episode; each draws a count in [min_turns, max_turns].""" + min_names_per_turn: int = 1 + """Minimum number of names introduced on each turn.""" + max_names_per_turn: int = 5 + """Maximum number of names introduced on each turn.""" + similarity_power: int = 4 + """Exponent applied to each turn's sequence-similarity score (higher = sharper penalty).""" + power_per_turn: bool = True + """Power-scale each turn then average (True), or average raw similarities then power once (False).""" + split: Literal["train"] = "train" + """Split of the source author-names dataset to build the episodes from.""" + + +class AlphabetSortTask(vf.Task): + info: dict + """The pre-generated episode: the `follow_ups` the user simulator reveals turn by turn, the + per-turn `ground_truths` the reward grades against, and `num_turns`.""" + + +class AlphabetSortTaskset(vf.Taskset[AlphabetSortTask, AlphabetSortConfig]): + def load_tasks(self) -> list[AlphabetSortTask]: + c = self.config + assert 1 <= c.min_turns <= c.max_turns, "need 1 <= min_turns <= max_turns" + assert 1 <= c.min_names_per_turn <= c.max_names_per_turn, ( + "need 1 <= min_names_per_turn <= max_names_per_turn" + ) + rng = random.Random(SEED) + tasks: list[AlphabetSortTask] = [] + for entry in load_dataset(DATASET, split=c.split): + names = list(dict.fromkeys(n.replace(" ", "") for n in entry["names"])) + counts = [ + rng.randint(c.min_names_per_turn, c.max_names_per_turn) + for _ in range(rng.randint(c.min_turns, c.max_turns)) + ] + if len(names) < sum(counts): + continue + by_first = rng.choice([True, False]) + label = "FIRST" if by_first else "LAST" + + def sort_key(s: str) -> str: + # split at the first capital after index 0 -> first- vs last-name part + cut = next((i for i in range(1, len(s)) if s[i].isupper()), len(s)) + return s[:cut] if by_first else s[cut:] + + turns, cumulative, ground_truths, i = [], [], [], 0 + for count in counts: + turn = names[i : i + count] + i += count + turns.append(turn) + cumulative += turn + ranked = sorted(cumulative, key=sort_key) + ground_truths.append( + ranked + if len(turns) == 1 + else [f"{x} // new name!" if x in turn else x for x in ranked] + ) + + first = turns[0][:] + rng.shuffle(first) + shown = rng.randint(c.min_names_per_turn, c.max_names_per_turn) + instruction = ( + f"Sort these names in alphabetical order by {label} name: {', '.join(first)}\n\n" + "Use exactly this format:\n\n" + + "\n".join(f"Name{j}" for j in range(1, shown + 1)) + + "\n" + ) + + follow_ups = [] + for t in range(1, len(turns)): + shuffled = turns[t][:] + rng.shuffle(shuffled) + shown = rng.randint( + c.min_names_per_turn, sum(len(x) for x in turns[: t + 1]) + ) + threshold = rng.randint(0, shown - 1) + prompt = ( + f"Now sort ALL of these names alphabetically by {label} name: {', '.join(shuffled)}\n\n" + "These are in addition to the prior list. Mark any NEW names (that weren't " + "in the prior list) with `// new name!` at the end." + ) + if t == 1: + prompt += ( + "\n\nUse exactly this format:\n\n" + + "\n".join( + f"Name{j}" + (" // new name!" if j > threshold else "") + for j in range(1, shown + 1) + ) + + "\n" + ) + else: + prompt += " Follow the same format as before." + follow_ups.append(prompt) + + tasks.append( + AlphabetSortTask( + idx=len(tasks), + instruction=instruction, + info={ + "follow_ups": follow_ups, + "ground_truths": ground_truths, + "num_turns": len(turns), + }, + ) + ) + return tasks + + def user(self, task: AlphabetSortTask) -> vf.User: + info = { + "follow_ups": task.info["follow_ups"], + "num_turns": task.info["num_turns"], + } + return vf.User( + name="user", + command=[sys.executable, "-m", "alphabet_sort_v1.user"], + env={"ALPHABET_SORT_INFO": json.dumps(info)}, + ) + + @vf.reward(weight=1.0) + async def alphabet_sort(self, task: AlphabetSortTask, trace: vf.Trace) -> float: + c = self.config + ground_truths = task.info["ground_truths"] + num_turns = task.info["num_turns"] + responses = [m.content or "" for m in trace.assistant_messages] + scores = [] + for t in range(num_turns): + tag = "alphabetical_sorted" if t == 0 else "combined_alphabetical_sorted" + response = responses[t] if t < len(responses) else "" + expected = "\n".join(s.strip().lower() for s in ground_truths[t]) + # Multiple attempts only count if they strictly improve (else 0). + attempts = [] + for content in re.findall(f"<{tag}>(.*?)", response, re.DOTALL): + pred = "\n".join( + ln.strip().lower() for ln in content.split("\n") if ln.strip() + ) + sim = ( + difflib.SequenceMatcher(None, pred, expected).ratio() + if pred and expected + else 0.0 + ) + attempts.append(sim**c.similarity_power if c.power_per_turn else sim) + if not attempts: + scores.append(0.0) + elif len(attempts) == 1: + scores.append(attempts[0]) + else: + improved = all(b > a for a, b in zip(attempts, attempts[1:])) + scores.append(attempts[-1] if improved else 0.0) + avg = sum(scores) / num_turns if num_turns else 0.0 + return avg if c.power_per_turn else avg**c.similarity_power + + +def load_taskset(config: AlphabetSortConfig) -> AlphabetSortTaskset: + return AlphabetSortTaskset(config) diff --git a/examples/tasksets/alphabet_sort_v1/alphabet_sort_v1/user.py b/examples/tasksets/alphabet_sort_v1/alphabet_sort_v1/user.py new file mode 100644 index 0000000000..bdbe135abd --- /dev/null +++ b/examples/tasksets/alphabet_sort_v1/alphabet_sort_v1/user.py @@ -0,0 +1,42 @@ +"""alphabet-sort-v1 user simulator: replays the episode's pre-generated follow-up turns. + +Launched by the framework as a per-rollout subprocess (a `vf.User`). It holds the task's +follow-up prompts and turn count (from `ALPHABET_SORT_INFO`) and serves a single `respond` +tool: after each assistant turn it injects the next follow-up as a user message, until all +turns are done. This colocates v0's `MultiTurnEnv.env_response` with the agent — the harness +drives it, never the model. +""" + +import json +import os + +from mcp.server.fastmcp import FastMCP + +import verifiers.v1 as vf + +INFO = json.loads(os.environ["ALPHABET_SORT_INFO"]) +FOLLOW_UPS = INFO["follow_ups"] +NUM_TURNS = INFO["num_turns"] + +mcp = FastMCP("user") + +# One `respond` call per assistant turn; track how many turns the model has taken. +_turns = 0 + + +@mcp.tool() +def respond(message: str) -> str: + """Inject the next follow-up turn, or end the episode once all turns are done.""" + global _turns + _turns += 1 + if _turns >= NUM_TURNS: + return json.dumps({"messages": [], "done": True}) + return json.dumps( + { + "messages": [{"role": "user", "content": FOLLOW_UPS[_turns - 1]}], + "done": False, + } + ) + + +vf.run_mcp_server(mcp) diff --git a/examples/tasksets/alphabet_sort_v1/pyproject.toml b/examples/tasksets/alphabet_sort_v1/pyproject.toml new file mode 100644 index 0000000000..2d7357fc5c --- /dev/null +++ b/examples/tasksets/alphabet_sort_v1/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "alphabet-sort-v1" +version = "0.1.0" +description = "alphabet-sort-v1 — maintain an alphabetically sorted list of names across turns, driven by a colocated user simulator." +requires-python = ">=3.10" +dependencies = ["datasets"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["alphabet_sort_v1"] diff --git a/pyproject.toml b/pyproject.toml index f286e66b86..fbbda0ab1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,7 +91,7 @@ examples = [ "compact", "gsm8k-v1", "wikispeedia-v1", "glossary-v1", "deepwiki-v1", "wiki-search-v1", "code-golf-v1", "reverse-text-v1", "hello-rlm-v1", "math-env-v1", "aime24-v1", - "wordle-v1", "terminal-bench-2-v1", + "wordle-v1", "terminal-bench-2-v1", "alphabet-sort-v1", ] [project.optional-dependencies] @@ -152,6 +152,7 @@ wordle-v1 = { path = "examples/tasksets/wordle_v1", editable = true } terminal-bench-2-v1 = { path = "examples/tasksets/terminal_bench_2_v1", editable = true } math-env-v1 = { path = "examples/tasksets/math_env_v1", editable = true } aime24-v1 = { path = "examples/tasksets/aime24_v1", editable = true } +alphabet-sort-v1 = { path = "examples/tasksets/alphabet_sort_v1", editable = true } [tool.uv.exclude-newer-package] # PrimeIntellect-published on PyPI (trusted publisher)