Adds Zebra/Murdle/Einstein/Grid Style Puzzles

README.md

@@ -115,6 +115,8 @@ See the [Dataset Gallery](GALLERY.md) for a complete list of available datasets
 - `PropositionalLogicDataset`: Generate propositional logic reasoning problems
 - `SyllogismDataset`: Generates a [syllogism](https://en.wikipedia.org/wiki/Syllogism) reasoning dataset
 - `AliceInWonderlandDataset`: Generates [AIW](https://openreview.net/forum?id=Mkl7dzjYiW) (Alice In Wonderland) problems with a few variations
+- `ZebraDataset`: Generates [Zebra Puzzles](https://en.wikipedia.org/wiki/Zebra_Puzzle) of varying difficulty.
+
 ### <small>Graph Tasks</small>
 
 - `FamilyRelationshipsDataset`: Generate family relationship reasoning tasks with family trees

pyproject.toml

@@ -16,8 +16,10 @@ dependencies = [
   "cellpylib==2.4.0",
   "sympy>=1.13.1",
   "magiccube==0.3.0",
+  "pycosat==0.6.6",
   "pyfiglet==1.0.2",
-  "pytz>=2024.1"
+  "pytz>=2024.1",
+  "tabulate==0.9.0",
 ]
 classifiers = [
   "Programming Language :: Python :: 3",

reasoning_gym/__init__.py

@@ -2,14 +2,15 @@
 Reasoning Gym - A library of procedural dataset generators for training reasoning models
 """
 
-from . import algebra, algorithmic, arithmetic, cognition, data, games, geometry, graphs, logic
+from . import algebra, algorithmic, arithmetic, code, cognition, data, games, geometry, graphs, logic
 from .factory import create_dataset, register_dataset
 
 __version__ = "0.1.3"
 __all__ = [
     "algebra",
     "algorithmic",
     "arithmetic",
+    "code",
     "cognition",
     "data",
     "games",

reasoning_gym/logic/__init__.py

@@ -9,6 +9,7 @@
 from .aiw import AliceInWonderlandConfig, AliceInWonderlandDataset
 from .propositional_logic import PropositionalLogicConfig, PropositionalLogicDataset
 from .syllogisms import SyllogismConfig, SyllogismDataset, Term
+from .zebra_puzzles import ZebraConfig, ZebraDataset
 
 __all__ = [
     "AliceInWonderlandConfig",
@@ -19,4 +20,6 @@
     "SyllogismDataset",
     "syllogism_dataset",
     "Term",
+    "ZebraConfig",
+    "ZebraDataset",
 ]

reasoning_gym/logic/contrib/logic_puzzle/README.md

@@ -0,0 +1,8 @@
+via https://github.com/nouhadziri/faith-and-fate
+
+@article{dziri2023faith,
+title={Faith and Fate: Limits of Transformers on Compositionality},
+author={Dziri, Nouha and Lu, Ximing and Sclar, Melanie and Li, Xiang Lorraine and Jian, Liwei and Lin, Bill Yuchen and West, Peter and Bhagavatula, Chandra and Bras, Ronan Le and Hwang, Jena D and others},
+journal={arXiv preprint arXiv:2305.18654},
+year={2023}
+}

reasoning_gym/logic/contrib/logic_puzzle/clues.py

@@ -0,0 +1,266 @@
+"""
+clues.py
+
+These are all the clue types that a puzzle can have. Things like "the tea drinker lives in the
+green house" and "the cat owner lives left of the person who likes grilled cheese."
+
+There's a Clue ABC that requires you implement an `as_cnf` method, to convert the clue to an
+and-of-ors (probably using things defined in `sat_utils`), and a human-readable __repr__ that
+can be used in a puzzle description.
+
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from functools import wraps
+from itertools import product
+from typing import Iterable, List, Tuple
+
+from reasoning_gym.logic.contrib.logic_puzzle.literals import Literal
+from reasoning_gym.logic.contrib.logic_puzzle.sat_utils import from_dnf, neg
+
+
+def _capitalize_first(repr_func):
+    """
+    Decorator for a __repr__ function that capitalizes the first letter without chagning the rest
+
+    (in contrast to str.capitalize(), which capitalizes the first letter and makes the rest lower)
+    """
+
+    @wraps(repr_func)
+    def wrapper(*args, **kwargs):
+        output = repr_func(*args, **kwargs)
+        return output[0].upper() + output[1:]
+
+    return wrapper
+
+
+class Clue(ABC):
+    """Base class for the types of clues that we allow."""
+
+    @abstractmethod
+    def as_cnf(self) -> Iterable[Tuple[str]]: ...
+
+    @abstractmethod
+    def __repr__(self) -> str: ...
+
+
+def comb(value: Literal, house: int) -> str:
+    """Format how a value is shown at a given house"""
+
+    return f"{value} {house}"
+
+
+@dataclass(eq=True, frozen=True)
+class found_at(Clue):
+    """
+    A literal is known to be at a specific house
+
+    Examples:
+     - the tea drinker lives in the middle house
+     - the fourth house is red
+    """
+
+    value: Literal
+    house: int
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return [(comb(self.value, self.house),)]
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        houses = [None, "first", "second", "third", "fourth", "fifth", "sixth", "seventh"]
+        return f"{self.value.value} is in the {houses[self.house]} house."
+
+
+@dataclass(eq=True, frozen=True)
+class not_at(Clue):
+    """
+    Two values are known *not* to be at the same house
+
+    Examples:
+     - the musician does not drink tea
+     - the red house does not contain a cat
+    """
+
+    value: Literal
+    house: int
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return [(neg(comb(self.value, self.house)),)]
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        houses = [None, "first", "second", "third", "fourth", "fifth", "sixth", "seventh"]
+        return f"{self.value.value} is not in the {houses[self.house]} house."
+
+
+@dataclass(eq=True, frozen=True)
+class same_house(Clue):
+    """
+    Two values are known to be at the same house
+
+    Examples:
+     - the musician drinks tea
+     - the red house contains a cat
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return from_dnf((comb(self.value1, i), comb(self.value2, i)) for i in self.houses)
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        return f"{self.value1.value} is {self.value2.value}."
+
+
+@dataclass(eq=True, frozen=True)
+class consecutive(Clue):
+    """
+    The first value is directly to the left of the second value
+
+    Examples:
+     - the green house is directly to the left of the white house
+       (green in 1, white in 2 OR green in 2, white in 3 OR etc.)
+     - the house with the kittens is directly to the right of the tea drinker's home
+       (kittens in 2, tea in 1 OR kittens in 3, tea in 2 OR etc.)
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return from_dnf((comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[1:]))
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        return f"{self.value1.value} is directly left of {self.value2.value}."
+
+
+@dataclass(eq=True, frozen=True)
+class beside(Clue):
+    """
+    The two values occur side-by-side (either left or right)
+
+    Examples:
+     - the coffee drinker is (left or right) of the tea drinker
+     - the cat owner is (left or right) of the green house
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return from_dnf(
+            [(comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[1:])]
+            + [(comb(self.value2, i), comb(self.value1, j)) for i, j in zip(self.houses, self.houses[1:])]
+        )
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        return f"{self.value1.value} and {self.value2.value} are next to each other."
+
+
+@dataclass(eq=True, frozen=True)
+class left_of(Clue):
+    """
+    The first value is somewhere to the left of the second value
+
+    Examples:
+     - the tea drinker is in house 1 and the musician in 2, 3, 4, or 5;
+       OR the tea drinker in 2, and musician in 3, 4, or 5;
+       OR the tea drinker in 3, musician in 4, 5; OR tea 4, musician 5.
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return from_dnf(
+            (comb(self.value1, i), comb(self.value2, j)) for i, j in product(self.houses, self.houses) if i < j
+        )
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        return f"{self.value1.value} is somewhere to the left of {self.value2.value}."
+
+
+@dataclass(eq=True, frozen=True)
+class right_of(Clue):
+    """
+    The first value is somewhere to the right of the second value.
+
+    Examples:
+     - the coffee drinker is in house 5 and the artist in 1, 2, 3, 4;
+       OR the coffee drinker in 4, and artist in 1, 2, or 3;
+       OR the coffee drinker in 3, artist in 1, 2; OR coffee 2, artist 1.
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return sat_utils.from_dnf(
+            (comb(self.value1, i), comb(self.value2, j)) for i, j in product(self.houses, self.houses) if i > j
+        )
+
+    @_capitalize_first
+    def __repr__(self) -> str:
+        return f"{self.value1.value} is somewhere to the right of {self.value2.value}."
+
+
+@dataclass(eq=True, frozen=True)
+class one_between(Clue):
+    """
+    The values are separated by one house
+
+    Examples (if 5 houses):
+     - the cat is in house 1 and tea drinker in house 3; OR cat 2, tea 4;
+       OR cat 4 house 5
+     - the green house is #1 and the musician in house 3; or green house 2, musician 4;
+       OR green house 3, musician 5.
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return from_dnf(
+            [(comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[2:])]
+            + [(comb(self.value2, i), comb(self.value1, j)) for i, j in zip(self.houses, self.houses[2:])]
+        )
+
+    def __repr__(self) -> str:
+        return f"There is one house between {self.value1.value} and {self.value2.value}."
+
+
+@dataclass(eq=True, frozen=True)
+class two_between(Clue):
+    """
+    The values are separated by two houses
+
+    Examples (if 5 houses):
+     - the cat is in house 1 and artist in house 4; or cat 2, artist 5
+     - the dog is in house 1 and red house is #4; or dog 2, red house 5
+    """
+
+    value1: Literal
+    value2: Literal
+    houses: Tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
+
+    def as_cnf(self) -> List[Tuple[str]]:
+        return from_dnf(
+            [(comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[3:])]
+            + [(comb(self.value2, i), comb(self.value1, j)) for i, j in zip(self.houses, self.houses[3:])]
+        )
+
+    def __repr__(self) -> str:
+        return f"There are two houses between {self.value1.value} and {self.value2.value}."