Files
logic-skeleton-bench/logicbench/skeletons.py
T

638 lines
22 KiB
Python

from __future__ import annotations
import random
from collections import deque
from dataclasses import dataclass, field
from typing import Any
from .models import Entity, HiddenProgram, Query, TransitionRule, WorldEvent, WorldSkeleton
LENGTH_PRESETS = {"short": 4, "medium": 6, "long": 8}
ANSWER_OPTIONS = ("YES", "NO")
@dataclass
class _World:
locations: dict[str, str]
order_level: dict[str, int]
set_level: dict[str, int]
above: dict[tuple[str, str], str] = field(default_factory=dict)
subsets: dict[tuple[str, str], str] = field(default_factory=dict)
members: dict[tuple[str, str], str] = field(default_factory=dict)
disjoint: dict[frozenset[str], str] = field(default_factory=dict)
connected: set[frozenset[str]] = field(default_factory=set)
numbers: dict[str, int] = field(default_factory=dict)
flags: dict[tuple[str, str], bool] = field(default_factory=dict)
def resolve_length(value: str | int) -> int:
"""Resolve required proof length, independently of event-log or text length."""
if isinstance(value, int):
result = value
elif value in LENGTH_PRESETS:
result = LENGTH_PRESETS[value]
else:
try:
result = int(value)
except ValueError as exc:
raise ValueError(f"length must be an integer or one of: {', '.join(LENGTH_PRESETS)}") from exc
if not 3 <= result <= 10:
raise ValueError("length must be between 3 and 10")
return result
def _program() -> HiddenProgram:
return HiddenProgram(
types=("Entity", "Set", "Region", "Number", "Boolean", "State"),
relations=(
"EQUAL",
"NOT_EQUAL",
"ABOVE",
"BELOW",
"MEMBER",
"SUBSET",
"DISJOINT",
"ADJACENT",
"CONNECTED",
"CAUSES",
"LOCATED_IN",
),
operators=(
"NOT",
"AND",
"OR",
"IMPLIES",
"ADD",
"SUBTRACT",
"COMPARE",
"UNION",
"INTERSECTION",
"MOVE",
"SWAP",
"SET",
"COMPOSE",
),
quantifiers=("ALL", "EXISTS"),
structures=("CHAIN", "BRANCH", "MERGE", "CYCLE", "NEST", "CONDITION", "SEQUENCE"),
inference_rules=(
"ABOVE(x,y) AND ABOVE(y,z) IMPLIES ABOVE(x,z)",
"ABOVE(x,y) IMPLIES BELOW(y,x)",
"MEMBER(x,A) AND SUBSET(A,B) IMPLIES MEMBER(x,B)",
"MEMBER(x,A) AND DISJOINT(A,B) IMPLIES NOT MEMBER(x,B)",
"CONNECTED is symmetric",
"Numeric updates compose in temporal order",
),
transition_rules=(
TransitionRule(
"move",
"SELECT_ENTITY",
("target region differs from current region",),
0.72,
("MOVE(entity, region)",),
),
TransitionRule(
"encounter-order",
"COLOCATED_PAIR",
("entities share a region", "latent order levels are adjacent"),
0.78,
("SET ABOVE(higher, lower)", "SET BELOW(lower, higher)"),
),
TransitionRule(
"membership",
"ENTITY_SET_CONTACT",
("set is at base level",),
0.68,
("SET MEMBER(entity, set)",),
),
TransitionRule(
"subset",
"SET_PAIR_CONTACT",
("set levels are adjacent",),
0.76,
("SET SUBSET(lower, higher)",),
),
TransitionRule(
"disjoint",
"SET_PAIR_CONTACT",
("sets have equal level", "sets are distinct"),
0.64,
("SET DISJOINT(left, right)",),
),
TransitionRule(
"numeric-update",
"COUNTER_PULSE",
(),
0.82,
("ADD or SUBTRACT(counter, amount)",),
),
TransitionRule(
"connect",
"ENTITY_PAIR_CONTACT",
(),
0.55,
("SET CONNECTED(left, right)",),
),
TransitionRule(
"state-change",
"ENTITY_SIGNAL",
(),
0.58,
("SET Boolean state",),
),
),
)
def _entities(entity_count: int) -> tuple[list[Entity], dict[str, list[str]]]:
if entity_count < 48:
raise ValueError("entity_count must be at least 48")
region_count = max(4, entity_count // 32)
set_count = max(24, entity_count // 5)
counter_count = max(8, entity_count // 16)
agent_count = entity_count - region_count - set_count - counter_count
groups = {
"Entity": [f"E{i:03d}" for i in range(agent_count)],
"Set": [f"S{i:03d}" for i in range(set_count)],
"Region": [f"R{i:03d}" for i in range(region_count)],
"Number": [f"N{i:03d}" for i in range(counter_count)],
}
entities = [Entity(identifier, kind) for kind, values in groups.items() for identifier in values]
return entities, groups
def _initial_world(
groups: dict[str, list[str]], rng: random.Random, proof_length: int
) -> tuple[_World, dict[str, Any]]:
agents, sets, regions, counters = (
groups["Entity"],
groups["Set"],
groups["Region"],
groups["Number"],
)
order_levels = max(proof_length + 2, 8)
set_levels = max(proof_length + 1, 8)
locations = {agent: rng.choice(regions) for agent in agents}
world = _World(
locations=locations,
order_level={agent: index % order_levels for index, agent in enumerate(agents)},
set_level={item: index % set_levels for index, item in enumerate(sets)},
numbers={counter: rng.randint(10, 90) for counter in counters},
)
public = [
{"id": f"I{index:04d}", "predicate": "LOCATED_IN", "arguments": [agent, region]}
for index, (agent, region) in enumerate(locations.items())
]
offset = len(public)
public.extend(
{
"id": f"I{offset + index:04d}",
"predicate": "INITIAL_VALUE",
"arguments": [counter, value],
}
for index, (counter, value) in enumerate(world.numbers.items())
)
hidden = {
"order_level": world.order_level,
"set_level": world.set_level,
"initial_locations": dict(locations),
"initial_numbers": dict(world.numbers),
}
return world, {"public_observations": public, "hidden_state": hidden}
def _event(
events: list[WorldEvent],
time: int,
kind: str,
participants: tuple[str, ...],
**arguments: Any,
) -> WorldEvent:
item = WorldEvent(f"V{len(events):05d}", time, kind, participants, arguments)
events.append(item)
return item
def _simulate(
world: _World,
groups: dict[str, list[str]],
steps: int,
rng: random.Random,
) -> list[WorldEvent]:
agents, sets, regions, counters = (
groups["Entity"],
groups["Set"],
groups["Region"],
groups["Number"],
)
events: list[WorldEvent] = []
kinds = ("ENCOUNTER", "MOVE", "MEMBER", "SUBSET", "DISJOINT", "NUMBER", "CONNECT", "STATE")
weights = (30, 12, 11, 15, 7, 11, 8, 6)
probabilities = {
"ENCOUNTER": 0.78,
"MOVE": 0.72,
"MEMBER": 0.68,
"SUBSET": 0.76,
"DISJOINT": 0.64,
"NUMBER": 0.82,
"CONNECT": 0.55,
"STATE": 0.58,
}
for time in range(1, steps + 1):
kind = rng.choices(kinds, weights=weights, k=1)[0]
if rng.random() >= probabilities[kind]:
continue
if kind == "MOVE":
agent = rng.choice(agents)
choices = [region for region in regions if region != world.locations[agent]]
destination = rng.choice(choices)
origin = world.locations[agent]
world.locations[agent] = destination
_event(events, time, "MOVE", (agent, origin, destination), origin=origin, destination=destination)
elif kind == "ENCOUNTER":
lower_level = rng.randrange(max(world.order_level.values()))
lower = [item for item in agents if world.order_level[item] == lower_level]
higher = [item for item in agents if world.order_level[item] == lower_level + 1]
if not lower or not higher:
continue
low, high = rng.choice(lower), rng.choice(higher)
if world.locations[low] != world.locations[high] or (high, low) in world.above:
continue
item = _event(
events,
time,
"RELATE",
(high, low),
predicate="ABOVE",
subject=high,
object=low,
region=world.locations[high],
)
world.above[(high, low)] = item.id
elif kind == "MEMBER":
agent = rng.choice(agents)
base_sets = [item for item in sets if world.set_level[item] == 0]
container = rng.choice(base_sets)
if (agent, container) in world.members:
continue
item = _event(
events,
time,
"RELATE",
(agent, container),
predicate="MEMBER",
subject=agent,
object=container,
)
world.members[(agent, container)] = item.id
elif kind == "SUBSET":
lower_level = rng.randrange(max(world.set_level.values()))
lower = [item for item in sets if world.set_level[item] == lower_level]
higher = [item for item in sets if world.set_level[item] == lower_level + 1]
child, parent = rng.choice(lower), rng.choice(higher)
if (child, parent) in world.subsets or any(left == child for left, _right in world.subsets):
continue
item = _event(
events,
time,
"RELATE",
(child, parent),
predicate="SUBSET",
subject=child,
object=parent,
)
world.subsets[(child, parent)] = item.id
elif kind == "DISJOINT":
level = rng.randrange(max(world.set_level.values()) + 1)
same_level = [item for item in sets if world.set_level[item] == level]
if len(same_level) < 2:
continue
left, right = rng.sample(same_level, 2)
key = frozenset((left, right))
if key in world.disjoint:
continue
item = _event(
events,
time,
"RELATE",
(left, right),
predicate="DISJOINT",
subject=left,
object=right,
)
world.disjoint[key] = item.id
elif kind == "NUMBER":
counter = rng.choice(counters)
delta = rng.choice((-7, -5, -3, 2, 4, 6))
before = world.numbers[counter]
after = max(0, before + delta)
world.numbers[counter] = after
_event(
events,
time,
"UPDATE_NUMBER",
(counter,),
operation="ADD" if delta >= 0 else "SUBTRACT",
amount=abs(delta),
before=before,
after=after,
)
elif kind == "CONNECT":
left, right = rng.sample(agents, 2)
key = frozenset((left, right))
if key in world.connected:
continue
world.connected.add(key)
_event(
events,
time,
"RELATE",
(left, right),
predicate="CONNECTED",
subject=left,
object=right,
)
else:
agent = rng.choice(agents)
flag = rng.choice(("ACTIVE", "SEALED", "VERIFIED"))
before = world.flags.get((agent, flag), False)
after = not before
world.flags[(agent, flag)] = after
_event(
events,
time,
"SET_STATE",
(agent,),
state=flag,
before=before,
after=after,
)
return events
def _shortest_path(
edges: dict[tuple[str, str], str], minimum: int
) -> tuple[str, str, tuple[str, ...], tuple[str, ...]] | None:
adjacency: dict[str, list[tuple[str, str]]] = {}
for (left, right), event_id in edges.items():
adjacency.setdefault(left, []).append((right, event_id))
candidates: list[tuple[str, str, tuple[str, ...], tuple[str, ...]]] = []
for source in adjacency:
queue = deque([(source, (source,), ())])
seen = {source}
while queue:
node, nodes, event_ids = queue.popleft()
if len(event_ids) >= minimum:
candidates.append((source, node, nodes, event_ids))
continue
for target, event_id in adjacency.get(node, []):
if target not in seen:
seen.add(target)
queue.append((target, nodes + (target,), event_ids + (event_id,)))
if not candidates:
return None
candidates.sort(key=lambda item: (len(item[3]), item[0], item[1]))
return candidates[0]
def _membership_path(
world: _World, minimum: int
) -> tuple[str, str, tuple[str, ...], tuple[str, ...]] | None:
subset_path = _shortest_path(world.subsets, minimum - 1)
if subset_path is None:
return None
candidates = []
for (entity, base_set), member_event in world.members.items():
path = _shortest_path_from(world.subsets, base_set, minimum - 1)
if path:
target, nodes, event_ids = path
candidates.append((entity, target, (member_event,) + event_ids, (entity,) + nodes))
if not candidates:
return None
candidates.sort(key=lambda item: (len(item[2]), item[0], item[1]))
entity, target, event_ids, nodes = candidates[0]
return entity, target, nodes, event_ids
def _membership_negative(
world: _World, minimum: int
) -> tuple[str, str, tuple[str, ...], tuple[str, ...]] | None:
adjacency: dict[str, list[tuple[str, str]]] = {}
for (left, right), event_id in world.subsets.items():
adjacency.setdefault(left, []).append((right, event_id))
for (entity, base_set), member_event in sorted(world.members.items()):
reachable = {base_set}
queue = deque([(base_set, (base_set,), ())])
paths: list[tuple[str, tuple[str, ...], tuple[str, ...]]] = []
while queue:
node, nodes, event_ids = queue.popleft()
paths.append((node, nodes, event_ids))
for target, event_id in adjacency.get(node, []):
if target not in reachable:
reachable.add(target)
queue.append((target, nodes + (target,), event_ids + (event_id,)))
for known_set, nodes, subset_events in paths:
if 1 + len(subset_events) < minimum - 1:
continue
for pair, disjoint_event in world.disjoint.items():
if known_set not in pair:
continue
other = next(item for item in pair if item != known_set)
if other in reachable:
continue
proof_ids = (member_event,) + subset_events + (disjoint_event,)
proof_nodes = (entity,) + nodes + (other,)
return entity, other, proof_nodes, proof_ids
return None
def _shortest_path_from(
edges: dict[tuple[str, str], str], source: str, minimum: int
) -> tuple[str, tuple[str, ...], tuple[str, ...]] | None:
adjacency: dict[str, list[tuple[str, str]]] = {}
for (left, right), event_id in edges.items():
adjacency.setdefault(left, []).append((right, event_id))
queue = deque([(source, (source,), ())])
seen = {source}
while queue:
node, nodes, event_ids = queue.popleft()
if len(event_ids) >= minimum:
return node, nodes, event_ids
for target, event_id in adjacency.get(node, []):
if target not in seen:
seen.add(target)
queue.append((target, nodes + (target,), event_ids + (event_id,)))
return None
def _choose_query(
index: int, complexity: int, world: _World, events: list[WorldEvent]
) -> tuple[Query, str, tuple[str, ...], tuple[str, ...], str] | None:
if index % 2 == 0:
path = _shortest_path(world.above, complexity)
if path is None:
return None
source, target, nodes, proof_ids = path
positive = index % 4 == 0
arguments = (source, target) if positive else (target, source)
gold = "YES" if positive else "NO"
facts = tuple(
f"ABOVE({nodes[position]},{nodes[position + 1]})"
for position in range(len(nodes) - 1)
)
if not positive:
facts += (f"ABOVE({source},{target}) therefore NOT ABOVE({target},{source})",)
return Query("ABOVE", arguments, events[-1].time), gold, proof_ids, facts, "spatial_order"
if index % 4 == 1:
negative = _membership_negative(world, complexity)
if negative is None:
return None
entity, target, nodes, proof_ids = negative
known_set, disjoint_set = nodes[-2], nodes[-1]
facts = [f"MEMBER({entity},{nodes[1]})"]
facts.extend(
f"SUBSET({nodes[position]},{nodes[position + 1]})"
for position in range(1, len(nodes) - 2)
)
facts.extend(
(
f"DISJOINT({known_set},{disjoint_set})",
f"therefore NOT MEMBER({entity},{target})",
)
)
return (
Query("MEMBER", (entity, target), events[-1].time),
"NO",
proof_ids,
tuple(facts),
"membership_disjointness",
)
membership = _membership_path(world, complexity)
if membership is None:
return None
entity, target, nodes, proof_ids = membership
facts = [f"MEMBER({entity},{nodes[1]})"]
facts.extend(
f"SUBSET({nodes[position]},{nodes[position + 1]})"
for position in range(1, len(nodes) - 1)
)
return (
Query("MEMBER", (entity, target), events[-1].time),
"YES",
proof_ids,
tuple(facts),
"membership_inheritance",
)
def _final_state(world: _World) -> dict[str, Any]:
return {
"locations": dict(world.locations),
"above": [[left, right] for left, right in sorted(world.above)],
"subsets": [[left, right] for left, right in sorted(world.subsets)],
"members": [[entity, container] for entity, container in sorted(world.members)],
"disjoint": [sorted(pair) for pair in sorted(world.disjoint, key=lambda item: sorted(item))],
"connected": [sorted(pair) for pair in sorted(world.connected, key=lambda item: sorted(item))],
"numbers": dict(world.numbers),
"flags": [
{"entity": entity, "state": state, "value": value}
for (entity, state), value in sorted(world.flags.items())
],
}
def _generate_one(
index: int,
complexity: int,
entity_count: int,
simulation_steps: int,
rng: random.Random,
) -> WorldSkeleton:
entities, groups = _entities(entity_count)
world, initial_state = _initial_world(groups, rng, complexity)
events = _simulate(world, groups, simulation_steps, rng)
if not events:
raise RuntimeError("simulation produced no events")
selected = _choose_query(index, complexity, world, events)
if selected is None:
raise RuntimeError("simulation did not produce a sufficiently deep proof")
query, gold, proof_ids, proof_facts, query_family = selected
direct = any(
event.kind == "RELATE"
and event.arguments.get("predicate") == query.predicate
and (event.arguments.get("subject"), event.arguments.get("object")) == query.arguments
for event in events
)
if direct:
raise RuntimeError("query was directly stated in the event history")
checks = {
"direct_query_event_absent": True,
"proof_length": len(proof_ids),
"minimum_proof_length": complexity,
"event_count": len(events),
"proof_event_ratio": len(proof_ids) / len(events),
"query_family": query_family,
"weak_baselines": {
"direct_fact": False,
"proof_shorter_than_requested": len(proof_ids) < complexity,
"proof_is_large_fraction_of_history": len(proof_ids) / len(events) >= 0.05,
},
}
return WorldSkeleton(
id=f"world-{index:04d}",
family="dynamic_multi_logic_world",
complexity=complexity,
entity_count=entity_count,
simulation_steps=simulation_steps,
hidden_program=_program(),
entities=tuple(entities),
initial_state=initial_state,
events=tuple(events),
final_state=_final_state(world),
query=query,
answer_options=ANSWER_OPTIONS,
gold_answer=gold,
proof_event_ids=proof_ids,
proof_facts=proof_facts,
generator_checks=checks,
)
def generate_skeletons(
count: int,
length: str | int,
seed: int = 0,
entity_count: int = 128,
simulation_steps: int = 1000,
) -> list[WorldSkeleton]:
if count < 1:
raise ValueError("count must be positive")
if simulation_steps < 100:
raise ValueError("simulation_steps must be at least 100")
complexity = resolve_length(length)
rng = random.Random(seed)
generated: list[WorldSkeleton] = []
for index in range(count):
for _attempt in range(60):
try:
generated.append(
_generate_one(index, complexity, entity_count, simulation_steps, rng)
)
break
except RuntimeError:
continue
else:
raise RuntimeError("could not generate a world with the requested proof complexity")
return generated