mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-29 17:35:16 +00:00
Add 13 new procedural datasets across 7 categories
New dataset categories: combinatorics, statistics, optimization, and formal languages. Extended existing algebra, arithmetic, probability, logic, and graphs packages with complex_advanced, linear_algebra, limits, number_theory, conditional_probability, set_operations, and job_scheduling. Each dataset includes config validation, deterministic seeding, custom scoring, curriculum support, and comprehensive unit tests (92 new tests).
This commit is contained in:
Ritvik19
parent
49b07130b3
commit
6eb252ae32
36 changed files with 3705 additions and 1 deletions
19
reasoning_gym/optimization/__init__.py
Normal file
19
reasoning_gym/optimization/__init__.py
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
"""
|
||||
Optimization reasoning tasks.
|
||||
"""
|
||||
|
||||
from .dynamic_programming import DynamicProgrammingConfig, DynamicProgrammingCurriculum, DynamicProgrammingDataset
|
||||
from .knapsack import KnapsackConfig, KnapsackCurriculum, KnapsackDataset
|
||||
from .linear_programming import LinearProgrammingConfig, LinearProgrammingCurriculum, LinearProgrammingDataset
|
||||
|
||||
__all__ = [
|
||||
"LinearProgrammingDataset",
|
||||
"LinearProgrammingConfig",
|
||||
"LinearProgrammingCurriculum",
|
||||
"KnapsackDataset",
|
||||
"KnapsackConfig",
|
||||
"KnapsackCurriculum",
|
||||
"DynamicProgrammingDataset",
|
||||
"DynamicProgrammingConfig",
|
||||
"DynamicProgrammingCurriculum",
|
||||
]
|
||||
199
reasoning_gym/optimization/dynamic_programming.py
Normal file
199
reasoning_gym/optimization/dynamic_programming.py
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..coaching import BaseCurriculum, ScalarAttributeDefinition
|
||||
from ..factory import ProceduralDataset, register_dataset
|
||||
|
||||
DATASET_NAME = "dynamic_programming"
|
||||
|
||||
TASK_TYPES = ("lcs", "coin_change", "lis", "edit_distance", "staircase")
|
||||
|
||||
|
||||
def _lcs_length(s1: str, s2: str) -> tuple[int, str]:
|
||||
m, n = len(s1), len(s2)
|
||||
dp = [[""] * (n + 1) for _ in range(m + 1)]
|
||||
for i in range(1, m + 1):
|
||||
for j in range(1, n + 1):
|
||||
if s1[i - 1] == s2[j - 1]:
|
||||
dp[i][j] = dp[i - 1][j - 1] + s1[i - 1]
|
||||
else:
|
||||
dp[i][j] = dp[i - 1][j] if len(dp[i - 1][j]) >= len(dp[i][j - 1]) else dp[i][j - 1]
|
||||
return len(dp[m][n]), dp[m][n]
|
||||
|
||||
|
||||
def _coin_change(coins: list[int], amount: int) -> int:
|
||||
dp = [float("inf")] * (amount + 1)
|
||||
dp[0] = 0
|
||||
for c in coins:
|
||||
for a in range(c, amount + 1):
|
||||
dp[a] = min(dp[a], dp[a - c] + 1)
|
||||
return dp[amount] if dp[amount] != float("inf") else -1
|
||||
|
||||
|
||||
def _lis_length(arr: list[int]) -> int:
|
||||
if not arr:
|
||||
return 0
|
||||
dp = [1] * len(arr)
|
||||
for i in range(1, len(arr)):
|
||||
for j in range(i):
|
||||
if arr[j] < arr[i]:
|
||||
dp[i] = max(dp[i], dp[j] + 1)
|
||||
return max(dp)
|
||||
|
||||
|
||||
def _edit_distance(s1: str, s2: str) -> int:
|
||||
m, n = len(s1), len(s2)
|
||||
dp = [[0] * (n + 1) for _ in range(m + 1)]
|
||||
for i in range(m + 1):
|
||||
dp[i][0] = i
|
||||
for j in range(n + 1):
|
||||
dp[0][j] = j
|
||||
for i in range(1, m + 1):
|
||||
for j in range(1, n + 1):
|
||||
if s1[i - 1] == s2[j - 1]:
|
||||
dp[i][j] = dp[i - 1][j - 1]
|
||||
else:
|
||||
dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
|
||||
return dp[m][n]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DynamicProgrammingConfig:
|
||||
min_str_len: int = 4
|
||||
max_str_len: int = 8
|
||||
min_arr_len: int = 5
|
||||
max_arr_len: int = 10
|
||||
max_staircase: int = 15
|
||||
task_types: tuple[str, ...] = TASK_TYPES
|
||||
task_weights: list[float] = field(default_factory=lambda: [0.2, 0.2, 0.2, 0.2, 0.2])
|
||||
seed: Optional[int] = None
|
||||
size: int = 500
|
||||
|
||||
def validate(self) -> None:
|
||||
assert self.size > 0, "size must be positive"
|
||||
assert self.min_str_len >= 2, "min_str_len must be >= 2"
|
||||
assert self.max_str_len >= self.min_str_len, "max_str_len must be >= min_str_len"
|
||||
assert self.min_arr_len >= 3, "min_arr_len must be >= 3"
|
||||
assert self.max_arr_len >= self.min_arr_len, "max_arr_len must be >= min_arr_len"
|
||||
assert self.max_staircase >= 2, "max_staircase must be >= 2"
|
||||
assert len(self.task_types) > 0, "must have at least one task type"
|
||||
assert all(t in TASK_TYPES for t in self.task_types), f"invalid task type"
|
||||
assert len(self.task_weights) == len(self.task_types), "weights must match types"
|
||||
|
||||
|
||||
class DynamicProgrammingDataset(ProceduralDataset):
|
||||
def __init__(self, config: DynamicProgrammingConfig):
|
||||
super().__init__(config=config, seed=config.seed, size=config.size)
|
||||
|
||||
def _make_lcs(self, rng: random.Random) -> dict:
|
||||
l1 = rng.randint(self.config.min_str_len, self.config.max_str_len)
|
||||
l2 = rng.randint(self.config.min_str_len, self.config.max_str_len)
|
||||
chars = "ABCDEFGH"
|
||||
s1 = "".join(rng.choice(chars) for _ in range(l1))
|
||||
s2 = "".join(rng.choice(chars) for _ in range(l2))
|
||||
length, _ = _lcs_length(s1, s2)
|
||||
question = (
|
||||
f"Find the length of the longest common subsequence (LCS) of '{s1}' and '{s2}'. "
|
||||
f"Give your answer as a single integer."
|
||||
)
|
||||
return {"question": question, "answer": str(length), "task_type": "lcs"}
|
||||
|
||||
def _make_coin_change(self, rng: random.Random) -> dict:
|
||||
num_coins = rng.randint(2, 4)
|
||||
coins = sorted(set(rng.randint(1, 10) for _ in range(num_coins + 2)))[:num_coins]
|
||||
if 1 not in coins:
|
||||
coins = [1] + coins
|
||||
amount = rng.randint(5, 25)
|
||||
result = _coin_change(coins, amount)
|
||||
question = (
|
||||
f"What is the minimum number of coins needed to make {amount} "
|
||||
f"using coins of denominations {coins}? Each denomination can be used unlimited times. "
|
||||
f"Give your answer as a single integer."
|
||||
)
|
||||
return {"question": question, "answer": str(result), "task_type": "coin_change"}
|
||||
|
||||
def _make_lis(self, rng: random.Random) -> dict:
|
||||
n = rng.randint(self.config.min_arr_len, self.config.max_arr_len)
|
||||
arr = [rng.randint(1, 50) for _ in range(n)]
|
||||
length = _lis_length(arr)
|
||||
question = (
|
||||
f"Find the length of the longest strictly increasing subsequence in {arr}. "
|
||||
f"Give your answer as a single integer."
|
||||
)
|
||||
return {"question": question, "answer": str(length), "task_type": "lis"}
|
||||
|
||||
def _make_edit_distance(self, rng: random.Random) -> dict:
|
||||
l1 = rng.randint(self.config.min_str_len, self.config.max_str_len)
|
||||
l2 = rng.randint(self.config.min_str_len, self.config.max_str_len)
|
||||
chars = "abcdefgh"
|
||||
s1 = "".join(rng.choice(chars) for _ in range(l1))
|
||||
s2 = "".join(rng.choice(chars) for _ in range(l2))
|
||||
dist = _edit_distance(s1, s2)
|
||||
question = (
|
||||
f"What is the minimum edit distance (Levenshtein distance) between "
|
||||
f"'{s1}' and '{s2}'? Operations: insert, delete, or substitute a character. "
|
||||
f"Give your answer as a single integer."
|
||||
)
|
||||
return {"question": question, "answer": str(dist), "task_type": "edit_distance"}
|
||||
|
||||
def _make_staircase(self, rng: random.Random) -> dict:
|
||||
n = rng.randint(3, self.config.max_staircase)
|
||||
ways = [0] * (n + 1)
|
||||
ways[0] = 1
|
||||
ways[1] = 1
|
||||
for i in range(2, n + 1):
|
||||
ways[i] = ways[i - 1] + ways[i - 2]
|
||||
question = (
|
||||
f"You are climbing a staircase with {n} steps. Each time you can climb 1 or 2 steps. "
|
||||
f"How many distinct ways can you reach the top? Give your answer as a single integer."
|
||||
)
|
||||
return {"question": question, "answer": str(ways[n]), "task_type": "staircase"}
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
rng = random.Random(self.seed + idx)
|
||||
task_type = rng.choices(self.config.task_types, weights=self.config.task_weights, k=1)[0]
|
||||
|
||||
generators = {
|
||||
"lcs": self._make_lcs,
|
||||
"coin_change": self._make_coin_change,
|
||||
"lis": self._make_lis,
|
||||
"edit_distance": self._make_edit_distance,
|
||||
"staircase": self._make_staircase,
|
||||
}
|
||||
result = generators[task_type](rng)
|
||||
return {
|
||||
"question": result["question"],
|
||||
"answer": result["answer"],
|
||||
"metadata": {
|
||||
"source_dataset": DATASET_NAME,
|
||||
"source_index": idx,
|
||||
"task_type": result["task_type"],
|
||||
"difficulty": {
|
||||
"max_str_len": self.config.max_str_len,
|
||||
"max_arr_len": self.config.max_arr_len,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class DynamicProgrammingCurriculum(BaseCurriculum):
|
||||
def __init__(self):
|
||||
super().__init__(DynamicProgrammingCurriculum.__name__, DynamicProgrammingConfig)
|
||||
self._define_attributes(
|
||||
ScalarAttributeDefinition(
|
||||
name="max_str_len",
|
||||
field_name="max_str_len",
|
||||
levels=[5, 8, 12, 15],
|
||||
description="Maximum string length",
|
||||
),
|
||||
ScalarAttributeDefinition(
|
||||
name="max_arr_len",
|
||||
field_name="max_arr_len",
|
||||
levels=[5, 10, 15, 20],
|
||||
description="Maximum array length",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
register_dataset(DATASET_NAME, DynamicProgrammingDataset, DynamicProgrammingConfig, DynamicProgrammingCurriculum)
|
||||
108
reasoning_gym/optimization/knapsack.py
Normal file
108
reasoning_gym/optimization/knapsack.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..coaching import BaseCurriculum, RangeAttributeDefinition, ScalarAttributeDefinition
|
||||
from ..factory import ProceduralDataset, register_dataset
|
||||
|
||||
DATASET_NAME = "knapsack"
|
||||
|
||||
|
||||
@dataclass
|
||||
class KnapsackConfig:
|
||||
min_items: int = 3
|
||||
max_items: int = 6
|
||||
min_weight: int = 1
|
||||
max_weight: int = 15
|
||||
min_value: int = 5
|
||||
max_value: int = 50
|
||||
min_capacity: int = 10
|
||||
max_capacity: int = 30
|
||||
seed: Optional[int] = None
|
||||
size: int = 500
|
||||
|
||||
def validate(self) -> None:
|
||||
assert self.size > 0, "size must be positive"
|
||||
assert self.min_items >= 2, "min_items must be >= 2"
|
||||
assert self.max_items >= self.min_items, "max_items must be >= min_items"
|
||||
assert self.min_weight >= 1, "min_weight must be >= 1"
|
||||
assert self.max_weight >= self.min_weight, "max_weight must be >= min_weight"
|
||||
assert self.min_value >= 1, "min_value must be >= 1"
|
||||
assert self.max_value >= self.min_value, "max_value must be >= min_value"
|
||||
assert self.min_capacity >= 1, "min_capacity must be >= 1"
|
||||
assert self.max_capacity >= self.min_capacity, "max_capacity must be >= min_capacity"
|
||||
|
||||
|
||||
def _solve_knapsack(weights: list[int], values: list[int], capacity: int) -> int:
|
||||
n = len(weights)
|
||||
dp = [[0] * (capacity + 1) for _ in range(n + 1)]
|
||||
for i in range(1, n + 1):
|
||||
for w in range(capacity + 1):
|
||||
dp[i][w] = dp[i - 1][w]
|
||||
if weights[i - 1] <= w:
|
||||
dp[i][w] = max(dp[i][w], dp[i - 1][w - weights[i - 1]] + values[i - 1])
|
||||
return dp[n][capacity]
|
||||
|
||||
|
||||
class KnapsackDataset(ProceduralDataset):
|
||||
def __init__(self, config: KnapsackConfig):
|
||||
super().__init__(config=config, seed=config.seed, size=config.size)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
rng = random.Random(self.seed + idx)
|
||||
|
||||
n = rng.randint(self.config.min_items, self.config.max_items)
|
||||
weights = [rng.randint(self.config.min_weight, self.config.max_weight) for _ in range(n)]
|
||||
values = [rng.randint(self.config.min_value, self.config.max_value) for _ in range(n)]
|
||||
capacity = rng.randint(self.config.min_capacity, self.config.max_capacity)
|
||||
|
||||
opt_val = _solve_knapsack(weights, values, capacity)
|
||||
|
||||
items_str = ", ".join(f"(weight={w}, value={v})" for w, v in zip(weights, values))
|
||||
question = (
|
||||
f"You have a knapsack with capacity {capacity}. "
|
||||
f"You have the following items: {items_str}. "
|
||||
f"Each item can be used at most once. "
|
||||
f"What is the maximum total value you can carry? "
|
||||
f"Give your answer as a single integer."
|
||||
)
|
||||
|
||||
return {
|
||||
"question": question,
|
||||
"answer": str(opt_val),
|
||||
"metadata": {
|
||||
"source_dataset": DATASET_NAME,
|
||||
"source_index": idx,
|
||||
"weights": weights,
|
||||
"values": values,
|
||||
"capacity": capacity,
|
||||
"difficulty": {
|
||||
"min_items": self.config.min_items,
|
||||
"max_items": self.config.max_items,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class KnapsackCurriculum(BaseCurriculum):
|
||||
def __init__(self):
|
||||
super().__init__(KnapsackCurriculum.__name__, KnapsackConfig)
|
||||
self._define_attributes(
|
||||
RangeAttributeDefinition(
|
||||
name="item_count",
|
||||
levels=[3, 6, 10, 15],
|
||||
lower_field_name="min_items",
|
||||
upper_field_name="max_items",
|
||||
description="Number of items",
|
||||
),
|
||||
RangeAttributeDefinition(
|
||||
name="capacity",
|
||||
levels=[10, 30, 50, 100],
|
||||
lower_field_name="min_capacity",
|
||||
upper_field_name="max_capacity",
|
||||
description="Knapsack capacity range",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
register_dataset(DATASET_NAME, KnapsackDataset, KnapsackConfig, KnapsackCurriculum)
|
||||
108
reasoning_gym/optimization/linear_programming.py
Normal file
108
reasoning_gym/optimization/linear_programming.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..coaching import BaseCurriculum, ScalarAttributeDefinition
|
||||
from ..factory import ProceduralDataset, register_dataset
|
||||
|
||||
DATASET_NAME = "linear_programming"
|
||||
|
||||
|
||||
@dataclass
|
||||
class LinearProgrammingConfig:
|
||||
min_coeff: int = 1
|
||||
max_coeff: int = 10
|
||||
num_constraints: int = 3
|
||||
seed: Optional[int] = None
|
||||
size: int = 500
|
||||
|
||||
def validate(self) -> None:
|
||||
assert self.size > 0, "size must be positive"
|
||||
assert self.min_coeff >= 1, "min_coeff must be >= 1"
|
||||
assert self.max_coeff >= self.min_coeff, "max_coeff must be >= min_coeff"
|
||||
assert 2 <= self.num_constraints <= 6, "num_constraints must be between 2 and 6"
|
||||
|
||||
|
||||
class LinearProgrammingDataset(ProceduralDataset):
|
||||
"""2-variable LP problems with backward construction from a known optimal vertex."""
|
||||
|
||||
def __init__(self, config: LinearProgrammingConfig):
|
||||
super().__init__(config=config, seed=config.seed, size=config.size)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
rng = random.Random(self.seed + idx)
|
||||
|
||||
x_opt = rng.randint(1, self.config.max_coeff)
|
||||
y_opt = rng.randint(1, self.config.max_coeff)
|
||||
c1 = rng.randint(self.config.min_coeff, self.config.max_coeff)
|
||||
c2 = rng.randint(self.config.min_coeff, self.config.max_coeff)
|
||||
opt_val = c1 * x_opt + c2 * y_opt
|
||||
|
||||
constraints = []
|
||||
for _ in range(self.config.num_constraints):
|
||||
a = rng.randint(self.config.min_coeff, self.config.max_coeff)
|
||||
b = rng.randint(self.config.min_coeff, self.config.max_coeff)
|
||||
rhs = a * x_opt + b * y_opt + rng.randint(0, 5)
|
||||
constraints.append((a, b, rhs))
|
||||
|
||||
tight_a = rng.randint(self.config.min_coeff, self.config.max_coeff)
|
||||
tight_b = rng.randint(self.config.min_coeff, self.config.max_coeff)
|
||||
tight_rhs = tight_a * x_opt + tight_b * y_opt
|
||||
constraints[0] = (tight_a, tight_b, tight_rhs)
|
||||
|
||||
constraint_strs = []
|
||||
for a, b, rhs in constraints:
|
||||
constraint_strs.append(f" {a}x + {b}y <= {rhs}")
|
||||
constraint_strs.append(" x >= 0")
|
||||
constraint_strs.append(" y >= 0")
|
||||
constraints_text = "\n".join(constraint_strs)
|
||||
|
||||
question = (
|
||||
f"Maximize {c1}x + {c2}y subject to:\n{constraints_text}\n"
|
||||
f"What is the maximum value of the objective function? "
|
||||
f"Give your answer as a single integer."
|
||||
)
|
||||
|
||||
return {
|
||||
"question": question,
|
||||
"answer": str(opt_val),
|
||||
"metadata": {
|
||||
"source_dataset": DATASET_NAME,
|
||||
"source_index": idx,
|
||||
"optimal_point": (x_opt, y_opt),
|
||||
"difficulty": {
|
||||
"num_constraints": self.config.num_constraints,
|
||||
"max_coeff": self.config.max_coeff,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
|
||||
if answer is None:
|
||||
return 0.0
|
||||
try:
|
||||
return 1.0 if int(answer.strip()) == int(entry["answer"]) else 0.0
|
||||
except (ValueError, TypeError):
|
||||
return 0.0
|
||||
|
||||
|
||||
class LinearProgrammingCurriculum(BaseCurriculum):
|
||||
def __init__(self):
|
||||
super().__init__(LinearProgrammingCurriculum.__name__, LinearProgrammingConfig)
|
||||
self._define_attributes(
|
||||
ScalarAttributeDefinition(
|
||||
name="num_constraints",
|
||||
field_name="num_constraints",
|
||||
levels=[2, 3, 4, 5],
|
||||
description="Number of inequality constraints",
|
||||
),
|
||||
ScalarAttributeDefinition(
|
||||
name="max_coeff",
|
||||
field_name="max_coeff",
|
||||
levels=[5, 10, 20, 50],
|
||||
description="Maximum coefficient value",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
register_dataset(DATASET_NAME, LinearProgrammingDataset, LinearProgrammingConfig, LinearProgrammingCurriculum)
|
||||
Loading…
Add table
Add a link
Reference in a new issue