feat(env): Course Schedule Curriculum (#266)

* course schedule curriculum

* update levels

* update comments

* lint
This commit is contained in:
Zafir Stojanovski 2025-03-05 22:42:46 +01:00 committed by GitHub
parent f3ee9a91a2
commit 8ccc4d7b0c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 104 additions and 29 deletions

View file

@ -10,10 +10,9 @@ from dataclasses import dataclass
from random import Random
from typing import Optional
from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
from ..factory import ProceduralDataset, register_dataset
MAX_NUM_COURSES = 1_000
QUESTION_TEMPLATE = """There are a total of {num_courses} courses you have to take, labeled from 0 to {last_index}.
You are given the following list of prerequisites, where prerequisites[i] = (a_i, b_i) indicates that you must first take course b_i if you want to take course a_i:
@ -27,25 +26,29 @@ Return True if you can finish all courses considering the prerequisites, or Fals
class CourseScheduleConfig:
"""Configuration for Course Schedule dataset generation"""
num_courses: int = 5 # Total number of courses (ranging from 0 to num_courses - 1)
min_num_courses: int = 5 # Minimum number of courses
max_num_courses: int = 10 # Maximum number of courses
min_num_prerequisites: int = 1 # Minimum number of prerequisites (per course)
max_num_prerequisites: int = 2 # Maximum number of prerequisites (per course)
p_solvable: float = 0.5 # Probability that the course schedule is solvable
min_cycle_length: int = 3 # Minimum length of a cycle in the prerequisites (if unsolvable)
max_cycle_length: int = 5 # Maximum length of a cycle in the prerequisites (if unsolvable)
p_solvable: float = 0.5 # Probability that the course schedule is solvable
size: int = 500 # Virtual dataset size
seed: Optional[int] = None
def validate(self):
"""Validate configuration parameters"""
assert 1 <= self.num_courses <= MAX_NUM_COURSES, f"num_courses must be between 1 and {MAX_NUM_COURSES}"
assert (
1 <= self.max_num_prerequisites <= self.num_courses
), "max_num_prerequisites must be between 1 and num_courses"
assert 0 <= self.p_solvable <= 1, "p_solvable must be between 0 and 1"
3 <= self.min_num_courses <= self.max_num_courses
), "min_num_courses must be between 3 and max_num_courses"
assert (
3 <= self.min_cycle_length <= self.max_cycle_length
), "min_cycle_length must be between 3 and max_cycle_length"
assert (
1 <= self.min_num_prerequisites <= self.max_num_prerequisites
), "min_num_prerequisites must be between 0 and max_num_prerequisites"
assert 0 <= self.p_solvable <= 1, "p_solvable must be between 0 and 1"
class CourseScheduleDataset(ProceduralDataset):
@ -88,14 +91,17 @@ class CourseScheduleDataset(ProceduralDataset):
for idx in range(len(courses) - 1, 0, -1):
current_course = courses[idx]
available_prereqs = courses[:idx] # Only earlier courses can be prerequisites
num_prerequisites = rng.randint(0, min(len(available_prereqs), self.config.max_num_prerequisites))
num_prerequisites = min(
len(available_prereqs),
rng.randint(self.config.min_num_prerequisites, self.config.max_num_prerequisites),
)
if num_prerequisites > 0:
chosen_prereqs = rng.sample(available_prereqs, num_prerequisites)
prerequisites.extend([[current_course, p] for p in chosen_prereqs])
if not solvable:
# If solution should be unsolvable, create a cycle
cycle_length = rng.randint(self.config.min_cycle_length, min(self.config.max_cycle_length, len(courses)))
cycle_length = min(len(courses), rng.randint(self.config.min_cycle_length, self.config.max_cycle_length))
cycle_courses = rng.sample(courses, cycle_length)
for i in range(cycle_length):
prerequisites.append([cycle_courses[i], cycle_courses[(i + 1) % cycle_length]])
@ -109,18 +115,19 @@ class CourseScheduleDataset(ProceduralDataset):
"""Generate a single Course Schedule question"""
rng = Random(self.seed + idx)
courses = list(range(self.config.num_courses))
num_courses = rng.randint(self.config.min_num_courses, self.config.max_num_courses)
courses = list(range(num_courses))
rng.shuffle(courses)
solvable = rng.random() < self.config.p_solvable
prerequisites = self._create_prerequisites(rng, courses, solvable)
answer = self._can_finish(self.config.num_courses, prerequisites)
answer = self._can_finish(num_courses, prerequisites)
return {
"question": QUESTION_TEMPLATE.format(
num_courses=self.config.num_courses,
last_index=self.config.num_courses - 1,
num_courses=num_courses,
last_index=num_courses - 1,
prerequisites=str(prerequisites),
),
"answer": str(answer),
@ -128,4 +135,43 @@ class CourseScheduleDataset(ProceduralDataset):
}
class CourseScheduleCurriculum(BaseCurriculum):
def __init__(self):
super().__init__(CourseScheduleCurriculum.__name__, CourseScheduleConfig)
# Define attributes
self._define_attributes(
RangeAttributeDefinition(
name="num_courses",
levels=[10, 50, 100, 500],
default_level=0, # Start with 5 courses
description="Number of courses in the schedule",
attr_type=AttributeType.APPEND,
min_value=3, # Ensure at least 3 courses
lower_field_name="min_num_courses",
upper_field_name="max_num_courses",
),
RangeAttributeDefinition(
name="num_prerequisites",
levels=[2, 3, 4, 5],
default_level=0, # Start with 2 prerequisites max
description="Number of prerequisites per course",
attr_type=AttributeType.APPEND,
min_value=0,
lower_field_name="min_num_prerequisites",
upper_field_name="max_num_prerequisites",
),
RangeAttributeDefinition(
name="cycle_length",
levels=[3, 4, 5, 6],
default_level=0, # Start with 3 cycle length
description="Length of a cycle in the prerequisites",
attr_type=AttributeType.APPEND,
min_value=3,
lower_field_name="min_cycle_length",
upper_field_name="max_cycle_length",
),
)
register_dataset("course_schedule", CourseScheduleDataset, CourseScheduleConfig)