Merge pull request #116 from zafstojano/env/pool_matrix

(Max/Average) Pool Matrix

reasoning_gym/algorithmic/__init__.py

@@ -19,6 +19,7 @@ from .manipulate_matrix import ManipulateMatrixConfig, ManipulateMatrixDataset
 from .number_filtering import NumberFilteringConfig, NumberFilteringDataset
 from .number_sorting import NumberSortingConfig, NumberSortingDataset
 from .palindrome_generation import PalindromeConfig, PalindromeDataset
+from .pool_matrix import PoolMatrixConfig, PoolMatrixDataset
 from .ransom_note import RansomNoteConfig, RansomNoteDataset
 from .rotate_matrix import RotateMatrixConfig, RotateMatrixDataset
 from .sentence_reordering import SentenceReorderingConfig, SentenceReorderingDataset
@@ -68,6 +69,8 @@ __all__ = [
     "ManipulateMatrixDataset",
     "BinaryMatrixConfig",
     "BinaryMatrixDataset",
+    "PoolMatrixConfig",
+    "PoolMatrixDataset",
     "ABConfig",
     "ABDataset",
     "CountPrimesConfig",

reasoning_gym/algorithmic/pool_matrix.py

@@ -0,0 +1,142 @@
+"""Perform average / max pooling on a matrix"""
+
+from copy import deepcopy
+from dataclasses import dataclass
+from random import Random
+from typing import Dict, Optional
+
+import numpy as np
+
+from ..factory import ProceduralDataset, register_dataset
+
+QUESTION_TEMPLATE = """Your job is to perform max/average pooling on the given matrix.
+The stride is equal to the kernel size, meaning there is no overlap between the pooling regions.
+
+Example 1:
+- Input: Perform max pooling on the following matrix with a kernel size of 2:
+1 2 3 4
+5 6 7 8
+9 10 11 12
+13 14 15 16
+- Output:
+6 8
+14 16
+
+Example 2:
+- Input: Perform average pooling on the following matrix with a kernel size of 2:
+1 2 3 4
+5 6 7 8
+9 10 11 12
+13 14 15 16
+- Output:
+3.5 5.5
+11.5 13.5
+
+Perform {pool_type} pooling on the following matrix with a kernel size of {pool_size}:
+{matrix}
+"""
+
+
+@dataclass
+class PoolMatrixConfig:
+    """Configuration for Pool Matrix dataset generation"""
+
+    min_rows: int = 2  # Minimum rows of the matrix
+    min_cols: int = 2  # Minimum columns of the matrix
+    max_rows: int = 10  # Maximum rows of the matrix
+    max_cols: int = 10  # Maximum columns of the matrix
+    max_pool_size: int = 3  # Maximum pooling size
+
+    size: int = 500  # Virtual dataset size
+    seed: Optional[int] = None
+
+    def validate(self):
+        """Validate configuration parameters"""
+        assert 2 <= self.min_rows, "min_rows must be at least 2"
+        assert 2 <= self.min_cols, "min_cols must be at least 2"
+        assert self.min_rows <= self.max_rows, "max_rows must be at least min_rows"
+        assert self.min_cols <= self.max_cols, "max_cols must be at least min_cols"
+        assert 1 <= self.max_pool_size, "max_pool_size must be at least 1"
+
+
+class PoolMatrixDataset(ProceduralDataset):
+    """Generates Pool Matrix exercises with configurable difficulty"""
+
+    def __init__(self, config: PoolMatrixConfig):
+        super().__init__(config=config, seed=config.seed, size=config.size)
+
+    def _get_matrix(self, rng: Random) -> np.ndarray:
+        """Generate a random matrix"""
+        rows = rng.randint(self.config.min_rows, self.config.max_rows)
+        cols = rng.randint(self.config.min_rows, self.config.max_cols)
+        return np.random.randint(0, 10, (rows, cols))
+
+    def _matrix_to_str(self, matrix: np.ndarray) -> str:
+        """Get a string representation of the matrix"""
+        return "\n".join(" ".join(str(round(x, 2)) for x in row) for row in matrix)
+
+    def _max_pool(self, matrix: np.ndarray, pool_size: int) -> np.ndarray:
+        """Perform max pooling on the matrix"""
+        rows, cols = matrix.shape
+        return np.array(
+            [
+                [np.max(matrix[i : i + pool_size, j : j + pool_size]) for j in range(0, cols, pool_size)]
+                for i in range(0, rows, pool_size)
+            ]
+        )
+
+    def _average_pool(self, matrix: np.ndarray, pool_size: int) -> np.ndarray:
+        """Perform average pooling on the matrix"""
+        rows, cols = matrix.shape
+        return np.array(
+            [
+                [np.mean(matrix[i : i + pool_size, j : j + pool_size]) for j in range(0, cols, pool_size)]
+                for i in range(0, rows, pool_size)
+            ]
+        )
+
+    def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float:
+        """Score the answer based on the metadata"""
+
+        reward = 0.0
+        try:
+            if answer is not None:
+                oracle_answer = np.array(entry["answer"])
+                answer = np.array(answer)
+                if oracle_answer.shape == answer.shape and np.allclose(oracle_answer, answer):
+                    reward = 1.0
+                if oracle_answer.shape == answer.shape:
+                    reward = 0.1
+                else:
+                    reward = 0.01
+        except:
+            print("Error in scoring answer for Pool Matrix")
+        return reward
+
+    def __getitem__(self, idx: int) -> dict:
+        """Generate a single Pool Matrix question"""
+        rng = Random(self.seed + idx)
+        np.random.seed(self.seed + idx)
+
+        matrix = self._get_matrix(rng)
+        matrix_str = self._matrix_to_str(matrix)
+
+        pool_size = rng.randint(1, self.config.max_pool_size)
+        pool_type = rng.choice(["average", "max"])
+
+        answer = self._average_pool(matrix, pool_size) if pool_type == "average" else self._max_pool(matrix, pool_size)
+        answer_str = self._matrix_to_str(answer)
+
+        return {
+            "question": QUESTION_TEMPLATE.format(matrix=matrix_str, pool_type=pool_type, pool_size=pool_size),
+            "answer": answer_str,
+            "metadata": {
+                "matrix": matrix.tolist(),
+                "pool_type": pool_type,
+                "pool_size": pool_size,
+                "solution": answer.tolist(),
+            },
+        }
+
+
+register_dataset("pool_matrix", PoolMatrixDataset, PoolMatrixConfig)

tests/test_pool_matrix.py

@@ -0,0 +1,138 @@
+"""Tests for Pool Matrix questions generation"""
+
+import numpy as np
+import pytest
+
+from reasoning_gym.algorithmic.pool_matrix import PoolMatrixConfig, PoolMatrixDataset
+
+
+def test_pool_matrix_config_validation():
+    """Test that invalid configs raise appropriate errors"""
+
+    for field in ["min_rows", "min_cols", "max_rows", "max_cols"]:
+        with pytest.raises(AssertionError):
+            config = PoolMatrixConfig(**{field: -1})  # Negative not allowed
+            config.validate()
+
+        with pytest.raises(AssertionError):
+            config = PoolMatrixConfig(**{field: 0})  # Zero not allowed
+            config.validate()
+
+        with pytest.raises(AssertionError):
+            config = PoolMatrixConfig(**{field: 1})  # One not allowed
+            config.validate()
+
+    with pytest.raises(AssertionError):
+        config = PoolMatrixConfig(max_pool_size=-1)  # Negative not allowed
+        config.validate()
+
+    with pytest.raises(AssertionError):
+        config = PoolMatrixConfig(max_pool_size=0)  # Zero not allowed
+        config.validate()
+
+
+def test_pool_matrix_dataset_deterministic():
+    """Test that dataset generates same items with same seed"""
+    config = PoolMatrixConfig(seed=42, size=10)
+    dataset1 = PoolMatrixDataset(config)
+    dataset2 = PoolMatrixDataset(config)
+
+    for i in range(len(dataset1)):
+        assert dataset1[i] == dataset2[i]
+
+
+def test_pool_matrix_dataset_items():
+    """Test basic properties of generated items"""
+    config = PoolMatrixConfig(max_rows=10, max_cols=10, max_pool_size=3, size=10, seed=42)
+    dataset = PoolMatrixDataset(config)
+
+    for i in range(len(dataset)):
+        item = dataset[i]
+        # Check item structure
+        assert isinstance(item, dict)
+        assert "question" in item
+        assert "answer" in item
+        assert "metadata" in item
+
+        # Check metadata
+        assert "matrix" in item["metadata"]
+        assert "pool_type" in item["metadata"]
+        assert "pool_size" in item["metadata"]
+        assert "solution" in item["metadata"]
+
+        matrix = item["metadata"]["matrix"]
+        pool_type = item["metadata"]["pool_type"]
+        pool_size = item["metadata"]["pool_size"]
+        solution = item["metadata"]["solution"]
+
+        # Verify dimensions
+        assert len(matrix) <= config.max_rows
+        assert all(len(row) <= config.max_cols for row in matrix)
+        assert len(solution) <= len(matrix)
+        assert len(solution[0]) <= len(matrix[0])
+        assert pool_size <= config.max_pool_size
+        assert pool_type in ["average", "max"]
+
+
+def test_pool_matrix_dataset_iteration():
+    """Test that iteration respects dataset size"""
+    config = PoolMatrixConfig(size=5, seed=42)
+    dataset = PoolMatrixDataset(config)
+
+    items = list(dataset)
+    assert len(items) == config.size
+
+    # Test multiple iterations yield same items
+    assert items == list(dataset)
+
+
+def test_pool_matrix_answer():
+    """Test the pooling methods"""
+    config = PoolMatrixConfig(seed=42)
+    dataset = PoolMatrixDataset(config)
+
+    # 1. Max pooling
+    matrix = np.array([[1, 2], [3, 4]])
+    assert np.allclose(dataset._max_pool(matrix, 2), np.array([[4]]))
+
+    matrix = np.array(
+        [
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+        ]
+    )
+    assert np.allclose(dataset._max_pool(matrix, 2), np.array([[6, 8], [10, 12]]))
+
+    matrix = np.array(
+        [
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+            [13, 14, 15, 16],
+        ]
+    )
+    assert np.allclose(dataset._max_pool(matrix, 2), np.array([[6, 8], [14, 16]]))
+
+    # 2. Average pooling
+    matrix = np.array([[1, 2], [3, 4]])
+    assert np.allclose(dataset._average_pool(matrix, 2), np.array([[2.5]]))
+
+    matrix = np.array(
+        [
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+        ]
+    )
+    assert np.allclose(dataset._average_pool(matrix, 2), np.array([[3.5, 5.5], [9.5, 11.5]]))
+
+    matrix = np.array(
+        [
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+            [13, 14, 15, 16],
+        ]
+    )
+    assert np.allclose(dataset._average_pool(matrix, 2), np.array([[3.5, 5.5], [11.5, 13.5]]))