fix(env): ARC 1D curriculum (#402)

* Add arc_1d curriculum

* Add difficulty to metadata

* use range attribute instead of scalar

---------

Co-authored-by: abdulhakeem <abdulhakeemadefioye@gmail.com>
Co-authored-by: Oliver Stanley <olivergestanley@gmail.com>
This commit is contained in:
Zafir Stojanovski 2025-04-01 13:01:15 +02:00 committed by GitHub
parent cd85c2d632
commit 50846c3534
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 69 additions and 3 deletions

View file

@ -2,6 +2,7 @@ from dataclasses import dataclass
from random import Random
from typing import Optional
from ..coaching import BaseCurriculum, RangeAttributeDefinition
from ..dataset import ProceduralDataset
from ..factory import register_dataset
@ -108,9 +109,31 @@ class Arc1DDataset(ProceduralDataset):
"size": size,
"train_examples": train_examples,
"test_example": test_example,
"difficulty": {
"size": (self.config.min_size, self.config.max_size),
},
},
}
class Arc1DCurriculum(BaseCurriculum):
"""Curriculum for ARC 1D tasks"""
def __init__(self):
super().__init__(Arc1DCurriculum.__name__, Arc1DConfig)
# Define attributes
self._define_attributes(
RangeAttributeDefinition(
name="size",
levels=[10, 25, 50, 100],
lower_field_name="min_size",
upper_field_name="max_size",
description="Grid size",
ensure_interval=True,
)
)
# Register the dataset
register_dataset(DATASET_NAME, Arc1DDataset, Arc1DConfig)
register_dataset(DATASET_NAME, Arc1DDataset, Arc1DConfig, Arc1DCurriculum)