fix(env): Unify CodeIO datasets (#405)

* unify codeio

* filtered for libraries not present in reasoning-gym

reasoning_gym/code/__init__.py

@@ -2,7 +2,14 @@
 Code reasing tasks
 """
 
-from .bf import BFConfig, BFDataset
-from .codeio import CodeIOConfig, CodeIODataset
+from .bf import BFConfig, BFCurriculum, BFDataset
+from .codeio import CodeIOConfig, CodeIOCurriculum, CodeIODataset
 
-__all__ = ["BFConfig", "BFDataset", "CodeIOConfig", "CodeIODataset"]
+__all__ = [
+    "BFConfig",
+    "BFDataset",
+    "BFCurriculum",
+    "CodeIOConfig",
+    "CodeIODataset",
+    "CodeIOCurriculum",
+]

reasoning_gym/code/codeio.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 
 import zss
 
+from ..coaching import BaseCurriculum, ScalarAttributeDefinition
 from ..data import get_data_file_path
 from ..factory import ProceduralDataset, register_dataset
 
@@ -59,10 +60,13 @@ class CodeIOConfig:
     seed: Optional[int] = None
     size: int = 500
     input_prediction_probability: float = 0.5
+    difficulty: Optional[int] = None
 
     def validate(self) -> None:
         """Validate configuration parameters"""
         assert 0.0 <= self.input_prediction_probability <= 1.0, "input_prediction_probability must be in [0, 1]"
+        if self.difficulty is not None:
+            assert 1 <= self.difficulty <= 10, "difficulty must be in [1, 10]"
 
 
 class CodeIODataset(ProceduralDataset):
@@ -80,18 +84,30 @@ class CodeIODataset(ProceduralDataset):
         self._data_path = get_data_file_path("codeio.jsonl.gz")
 
         with gzip.open(self._data_path, "rt", encoding="utf-8") as f:
-            CodeIODataset._jsonl_data = [json.loads(line) for line in f]
+            data = [json.loads(line) for line in f]
+            if self.config.difficulty is not None:
+                data = [entry for entry in data if entry.get("difficulty", -1) == self.config.difficulty]
+            assert len(data) > 0, "No data found for the specified difficulty level"
+            CodeIODataset._jsonl_data = data
+
+    def _generate_io_pair(self, main_code: str, input_generator_code: str, rng: Random, max_retries: int = 1):
+        local_vars = {"Random": Random}
+
+        full_code = f"{main_code}\n\n{input_generator_code}"
+        try:
+            exec(full_code, local_vars, local_vars)
+        except Exception as e:
+            print(f"Error executing code:\n{full_code}")
+            print(f"---------------------\nException: {e}\n---------------------")
+            return {}, {}
 
-    def _generate_io_pair(self, main_code: str, input_generator_code: str, rng: Random, max_retries: int = 3):
-        local_vars = {}
-        exec(main_code, {"Random": Random}, local_vars)
-        exec(input_generator_code, {"Random": Random}, local_vars)
         for _ in range(max_retries):
             try:
                 inputs = local_vars["generate_inputs"](rng)
                 outputs = local_vars["main_solution"](**inputs)
-            except Exception:
+            except Exception as e:
                 # Retry
+                print(f"Error generating I/O pair: {e}")
                 continue
             return inputs, outputs
         return {}, {}
@@ -124,6 +140,7 @@ class CodeIODataset(ProceduralDataset):
                 "source_index": idx,
                 "input_data": input_data,
                 "output_data": output_data,
+                "difficulty": {"difficulty": self.config.difficulty},
             },
         }
 
@@ -237,5 +254,19 @@ class CodeIODataset(ProceduralDataset):
         return reward
 
 
+class CodeIOCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(CodeIOCurriculum.__name__, CodeIOConfig)
+
+        self._define_attributes(
+            ScalarAttributeDefinition(
+                name="difficulty",
+                field_name="difficulty",
+                levels=[6, 7, 8, 9],
+                description="Difficulty level of the task",
+            ),
+        )
+
+
 # Register the dataset
-register_dataset(DATASET_NAME, CodeIODataset, CodeIOConfig)
+register_dataset(DATASET_NAME, CodeIODataset, CodeIOConfig, CodeIOCurriculum)

tests/test_codeio.py

@@ -1,6 +1,6 @@
 import pytest
 
-from reasoning_gym.code.codeio import CodeIOConfig, CodeIODataset
+from reasoning_gym.code.codeio import CodeIOConfig, CodeIOCurriculum, CodeIODataset
 
 
 def test_codeio_dataset():
@@ -40,3 +40,24 @@ def test_codeio_config():
 
     CodeIOConfig(size=10, seed=42, input_prediction_probability=0.1).validate()
     CodeIOConfig(size=10, seed=42, input_prediction_probability=0.9).validate()
+
+
+def test_codeio_curriculum():
+    curriculum = CodeIOCurriculum()
+
+    base_value = {"size": 150, "seed": 1}
+
+    base_cfg: CodeIOConfig = curriculum.generate_configuration(base_value)
+    assert base_cfg.seed == 1
+    assert base_cfg.size == 150
+    assert base_cfg.difficulty == 6
+
+    # test incrementing attribute level
+    curriculum.increment_attr_level("difficulty")
+    increased_cfg = curriculum.generate_configuration(base_value)
+    assert increased_cfg.difficulty == 7
+
+    # test decrementing attribute level
+    curriculum.decrement_attr_level("difficulty")
+    partially_decreased_cfg = curriculum.generate_configuration(base_value)
+    assert partially_decreased_cfg.difficulty == 6