""" Adaptive curriculum manager for dual-task math training. """ from __future__ import annotations import json import logging import random from dataclasses import asdict, dataclass, field from datetime import datetime from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np from src.rl.question_classifier import QuestionClassifier, TOPIC_LIST logger = logging.getLogger(__name__) # Maps dataset skill_id prefixes → canonical curriculum topic names. # Used during bootstrap so questions from NuminaMath / OpenMathInstruct # are credited to the right ZPD bucket instead of being misclassified. SKILL_ID_TO_TOPIC: dict[str, str] = { # NuminaMath-CoT "numina_algebra": "algebra", "numina_prealgebra": "algebra", "numina_number_theory": "number_theory", "numina_geometry": "geometry", "numina_combinatorics": "combinatorics", "numina_calculus": "calculus", "numina_statistics": "statistics", "numina_synthetic": "multi_step_reasoning", "numina_olympiad": "competition_math", "numina_competition": "competition_math", "numina_general": "multi_step_reasoning", # OpenMathInstruct-2 "openmath_algebra": "algebra", "openmath_prealgebra": "algebra", "openmath_number_theory":"number_theory", "openmath_geometry": "geometry", "openmath_combinatorics":"combinatorics", "openmath_calculus": "calculus", "openmath_competition": "competition_math", "openmath_synthetic": "multi_step_reasoning", "openmath_general": "multi_step_reasoning", # Legacy "gsm8k_grade_school": "basic_arithmetic", "aqua_rat_algebra": "algebra", "question_generation": "multi_step_reasoning", } @dataclass class TopicState: topic_name: str total_attempts: int = 0 successes: int = 0 success_rate: float = 0.0 difficulty_target: float = 0.5 difficulty_history: List[float] = field(default_factory=list) last_practiced: int = 0 first_attempted: int = 0 status: str = "untested" mastered_at_iteration: Optional[int] = None retention_tests_passed: int = 0 last_retention_score: float = 0.0 consecutive_failures: int = 0 failure_count_total: int = 0 history: List[Dict[str, float]] = field(default_factory=list) current_iteration_attempts: int = 0 # Track attempts within current iteration class CurriculumManager: """Goldilocks curriculum with adaptive topic selection.""" SWEET_SPOT_MIN = 0.4 SWEET_SPOT_MAX = 0.7 TARGET_SUCCESS = 0.55 CONTEXTS = ["bakery", "shopping", "school", "sports", "gardening", "travel"] ACTIONS = ["uses", "sells", "shares", "loses", "earns", "mixes"] TOPIC_TEMPLATES = { "basic_arithmetic": [ "Generate a {context} word problem using addition/subtraction.", ], "single_step_word_problems": [ "Generate a simple one-idea word problem in a {context} setting.", ], "fractions": [ "Generate a fractions word problem where someone {action} part of a quantity in a {context} scenario.", "Create a problem involving fraction operations in {context}.", ], "percentages": [ "Generate a percentage change or discount problem in {context}.", ], "ratios": [ "Generate a ratios/proportions word problem in {context}.", ], "money_problems": [ "Create a money and pricing problem in {context}.", ], "time_distance": [ "Generate a time/speed/distance problem in {context}.", ], "multi_step_reasoning": [ "Generate a multi-step reasoning problem in {context}.", ], "algebra": [ "Generate an algebra problem that solves for a variable in {context}.", ], "mixed_operations": [ "Generate a problem requiring mixed operations in {context}.", ], "comparison_problems": [ "Generate a comparison problem in {context} ('more than'/'less than').", ], "optimization_problems": [ "Generate a constrained optimization style word problem in {context}.", ], # ── AQuA-RAT additions ──────────────────────────────────────────── "number_theory": [ "Generate a number theory problem about divisibility, remainders, or prime factors in a {context} setting.", "Create a problem involving multiples and factors where someone in {context} {action} items in groups.", ], "profit_loss": [ "Generate a profit and loss problem where someone in {context} {action} goods at cost price and selling price.", "Create a problem about percentage profit or loss on a transaction in {context}.", ], "interest": [ "Generate a simple or compound interest problem involving a loan or investment in {context}.", "Create a problem where someone in {context} {action} money at a given annual interest rate.", ], "sets": [ "Generate a set theory or Venn diagram problem where people in {context} belong to overlapping groups.", "Create a problem using union and intersection of two groups in {context}.", ], "combinatorics": [ "Generate a combinatorics problem about arrangements or selections of objects in {context}.", "Create a problem involving permutations or combinations where someone in {context} {action} items.", ], "sequences": [ "Generate an arithmetic or geometric sequence problem in {context}.", "Create a problem where someone in {context} follows a pattern and must find the nth term.", ], "probability": [ "Generate a probability problem involving random selection or chance events in {context}.", "Create a problem where someone in {context} {action} items from a group and asks for probability.", ], "work_time": [ "Generate a work-rate problem where two people in {context} complete a task together or alone.", "Create a problem where workers in {context} {action} a job at different rates.", ], # ── NuminaMath / OpenMathInstruct additions ─────────────────────── "geometry": [ "Generate a geometry problem about area or perimeter of a shape encountered in {context}.", "Create a problem involving triangles or circles where someone in {context} needs to find a missing length or angle.", "Generate a coordinate geometry problem where points in a {context} layout form a geometric figure.", "Create a problem involving volume or surface area of a 3D shape relevant to {context}.", ], "calculus": [ "Generate a rate-of-change problem where a quantity in {context} grows or shrinks over time.", "Create an optimization problem where someone in {context} wants to maximise profit or minimise cost using calculus.", "Generate a problem involving a function whose minimum or maximum value must be found in a {context} scenario.", ], "statistics": [ "Generate a statistics problem where someone in {context} collects data and must find the mean, median, or mode.", "Create a problem involving standard deviation or variance of measurements taken in {context}.", "Generate a problem where data from {context} is summarised and an outlier or expected value must be identified.", ], "competition_math": [ "Generate a number theory problem asking how many positive integers satisfy a divisibility condition.", "Create a competition-style problem: find all integer solutions to an equation involving remainders or modular arithmetic.", "Generate a counting problem asking in how many ways objects can be arranged or selected under a constraint.", "Create a problem: given two integers relatively prime to each other, find their least common multiple or sum of digits.", ], } TOPIC_PREREQUISITES = { "fractions": ["basic_arithmetic"], "percentages": ["fractions"], "ratios": ["basic_arithmetic"], "algebra": ["basic_arithmetic", "comparison_problems"], "mixed_operations": ["basic_arithmetic", "fractions"], "optimization_problems": ["comparison_problems", "algebra"], # AQuA-RAT additions "number_theory": ["basic_arithmetic"], "profit_loss": ["percentages", "money_problems"], "interest": ["percentages"], "sets": ["basic_arithmetic"], "combinatorics": ["basic_arithmetic"], "sequences": ["basic_arithmetic", "algebra"], "probability": ["fractions", "ratios"], "work_time": ["ratios", "multi_step_reasoning"], # NuminaMath / OpenMathInstruct additions "geometry": ["basic_arithmetic"], "calculus": ["algebra", "sequences"], "statistics": ["ratios", "fractions"], "competition_math": ["number_theory", "combinatorics", "algebra"], } def __init__(self, checkpoint_dir: str | Path): self.checkpoint_dir = Path(checkpoint_dir) self.checkpoint_dir.mkdir(parents=True, exist_ok=True) self.classifier = QuestionClassifier() self.current_iteration = 0 self.recent_combined_rewards: List[float] = [] self.topics: Dict[str, TopicState] = { topic: TopicState(topic_name=topic) for topic in TOPIC_LIST } self.current_focus_topics: List[str] = [] self.hyperparams = { "sweet_spot_min": self.SWEET_SPOT_MIN, "sweet_spot_max": self.SWEET_SPOT_MAX, "target_success": self.TARGET_SUCCESS, } def initialize(self, bootstrap_questions: Optional[List[str]] = None) -> None: """Initialize topic priors, optionally from GSM8K-style questions.""" if bootstrap_questions: counts = {topic: 0 for topic in TOPIC_LIST} for question in bootstrap_questions: detected = self.classifier.classify_topic(question) topic = str(detected["primary_topic"]) if topic in counts: counts[topic] += 1 total = max(1, sum(counts.values())) for topic, state in self.topics.items(): prevalence = counts[topic] / total state.difficulty_target = float(max(0.3, min(0.75, 0.35 + 0.8 * prevalence))) else: for state in self.topics.values(): state.difficulty_target = 0.5 def initialize_from_dataset( self, records: List[Dict], difficulty_field: str = "difficulty", ) -> None: """ Bootstrap curriculum topic priors directly from dataset skill_ids. Much faster than the question-classifier bootstrap path — reads skill_id and difficulty from each JSONL record rather than running the keyword classifier on every question. Args: records: List of dataset records with 'skill_id' and optionally 'difficulty' fields. difficulty_field: Name of the difficulty field (default 'difficulty'). Values: 1=easy, 2=medium, 3=hard. """ counts: Dict[str, int] = {topic: 0 for topic in TOPIC_LIST} # Map difficulty 1/2/3 → difficulty_target 0.35/0.55/0.75 _diff_map = {1: 0.35, 2: 0.55, 3: 0.75} topic_difficulties: Dict[str, List[float]] = {t: [] for t in TOPIC_LIST} for rec in records: skill_id = rec.get("skill_id", "") topic = SKILL_ID_TO_TOPIC.get(skill_id) if topic is None: # Fall back to keyword classifier on the question text msgs = rec.get("messages", []) question = next( (m.get("content", "") for m in msgs if m.get("role") == "user"), "" ) if question: detected = self.classifier.classify_topic(question) topic = str(detected["primary_topic"]) else: continue if topic not in counts: continue counts[topic] += 1 raw_diff = rec.get(difficulty_field, 2) topic_difficulties[topic].append(_diff_map.get(int(raw_diff), 0.55)) total = max(1, sum(counts.values())) for topic, state in self.topics.items(): prevalence = counts[topic] / total # Difficulty target: average of observed difficulties, biased up # for rare topics (less data → start harder to find signal faster) diffs = topic_difficulties[topic] if diffs: avg_diff = sum(diffs) / len(diffs) else: avg_diff = 0.50 # Blend prevalence-based prior with observed average difficulty state.difficulty_target = float( max(0.30, min(0.80, 0.4 * avg_diff + 0.6 * (0.35 + 0.8 * prevalence))) ) logger.info( "Curriculum bootstrapped from %d records across %d topics", len(records), sum(1 for c in counts.values() if c > 0), ) for topic, cnt in sorted(counts.items(), key=lambda x: -x[1]): if cnt > 0: logger.debug( " %-30s %4d samples target_difficulty=%.2f", topic, cnt, self.topics[topic].difficulty_target, ) def select_topic_and_difficulty(self) -> Tuple[str, float]: probs = self._compute_topic_probabilities() names = list(probs.keys()) dist = np.array([probs[name] for name in names], dtype=np.float64) dist = dist / dist.sum() # Log topic distribution at start of each iteration (rollout 0, 10, 20, etc.) total_attempts = sum(t.current_iteration_attempts for t in self.topics.values()) if total_attempts % 20 == 0: # Every 20 rollouts top_5 = sorted(probs.items(), key=lambda x: x[1], reverse=True)[:5] logger.info(f"Topic probabilities (rollout {total_attempts}): {[(t, f'{p:.3f}') for t, p in top_5]}") topic = str(np.random.choice(names, p=dist)) difficulty = self._get_difficulty_for_topic(topic) self.current_focus_topics = [topic] return topic, difficulty def update_from_trajectory( self, topic: str, question_reward: float, solution_success: bool, combined_reward: Optional[float] = None, measured_difficulty: Optional[float] = None, ) -> None: # Grounded rollouts tag themselves with a synthetic topic name # (``grounded_gsm8k``) that isn't part of the curriculum ontology. # They must not pollute per-topic statistics — silently skip the # update for any unknown topic instead of crashing. The combined # reward is still recorded for plateau detection. state = self.topics.get(topic) if state is None: logger.debug( "Skipping curriculum update for out-of-ontology topic %r", topic ) if combined_reward is not None: self.recent_combined_rewards.append(float(combined_reward)) self.recent_combined_rewards = self.recent_combined_rewards[-30:] return state.total_attempts += 1 state.current_iteration_attempts += 1 state.successes += int(solution_success) state.success_rate = state.successes / max(1, state.total_attempts) state.last_practiced = self.current_iteration if state.first_attempted == 0: state.first_attempted = self.current_iteration success_value = 1.0 if solution_success else 0.0 if solution_success: state.consecutive_failures = 0 else: state.consecutive_failures += 1 state.failure_count_total += 1 target = state.difficulty_target # Only adjust difficulty if we have sufficient data if state.total_attempts >= 5: if state.success_rate > self.SWEET_SPOT_MAX: # Increase difficulty gradually state.difficulty_target = min(0.95, target + 0.03) if state.status != "mastered" and state.success_rate >= 0.75: state.status = "mastered" state.mastered_at_iteration = self.current_iteration elif state.success_rate < self.SWEET_SPOT_MIN: # Decrease difficulty more conservatively to avoid getting stuck too low state.difficulty_target = max(0.2, target - 0.04) state.status = "active" else: state.status = "active" if measured_difficulty is not None: state.difficulty_history.append(float(measured_difficulty)) else: state.difficulty_history.append(state.difficulty_target) state.history.append( { "iteration": float(self.current_iteration), "question_reward": float(question_reward), "solution_success": float(success_value), "success_rate": float(state.success_rate), "difficulty_target": float(state.difficulty_target), } ) if combined_reward is not None: self.recent_combined_rewards.append(float(combined_reward)) self.recent_combined_rewards = self.recent_combined_rewards[-30:] self.handle_persistent_failure(topic) def increment_iteration(self) -> None: self.current_iteration += 1 # Reset within-iteration counters for state in self.topics.values(): state.current_iteration_attempts = 0 self._run_retention_tests_if_due() def generate_instruction(self, topic: str, target_difficulty: float) -> str: templates = self.TOPIC_TEMPLATES.get(topic, self.TOPIC_TEMPLATES["multi_step_reasoning"]) template = random.choice(templates) # Note: {steps} placeholder removed from templates to let model decide complexity return template.format( context=random.choice(self.CONTEXTS), action=random.choice(self.ACTIONS), ) def get_curriculum_stats(self) -> Dict[str, object]: return { "iteration": self.current_iteration, "topics": {topic: asdict(state) for topic, state in self.topics.items()}, "sweet_spot_topics": self.get_sweet_spot_topics(), "current_focus_topics": self.get_current_focus(), "avg_recent_reward": float(np.mean(self.recent_combined_rewards)) if self.recent_combined_rewards else 0.0, } def get_sweet_spot_topics(self) -> List[str]: return [ topic for topic, state in self.topics.items() if state.total_attempts > 0 and self.SWEET_SPOT_MIN <= state.success_rate <= self.SWEET_SPOT_MAX ] def get_current_focus(self) -> List[str]: return list(self.current_focus_topics) def save_state(self, iteration: int, rollout: Optional[int] = None) -> None: if rollout is not None and rollout % 10 != 0: return filename = ( f"iteration_{iteration:03d}_final.json" if rollout is None else f"iteration_{iteration:03d}_rollout_{rollout:03d}.json" ) path = self.checkpoint_dir / filename state = { "version": "1.0", "timestamp": datetime.utcnow().isoformat(), "iteration": iteration, "rollout": rollout, "current_iteration": self.current_iteration, "recent_combined_rewards": self.recent_combined_rewards, "topics": {topic: asdict(topic_state) for topic, topic_state in self.topics.items()}, "hyperparams": self.hyperparams, } path.write_text(json.dumps(state, indent=2), encoding="utf-8") def load_checkpoint_safe(self) -> bool: checkpoints = sorted(self.checkpoint_dir.glob("iteration_*_final.json"), reverse=True) for checkpoint in checkpoints: try: data = json.loads(checkpoint.read_text(encoding="utf-8")) topics = data["topics"] if not isinstance(topics, dict): raise ValueError("Invalid topics section") rebuilt = {} for topic in TOPIC_LIST: values = topics.get(topic) if values is None or "success_rate" not in values: raise ValueError(f"Topic {topic} missing or malformed") rebuilt[topic] = TopicState(**values) self.topics = rebuilt self.current_iteration = int(data.get("current_iteration", data.get("iteration", 0))) self.recent_combined_rewards = list(data.get("recent_combined_rewards", [])) logger.info("Loaded curriculum state from %s", checkpoint) return True except Exception as exc: # pragma: no cover - defensive logger.warning("Failed to load curriculum checkpoint %s: %s", checkpoint, exc) return False def _compute_topic_probabilities(self) -> Dict[str, float]: all_states = list(self.topics.values()) sweet_spot = [t for t in all_states if self.SWEET_SPOT_MIN <= t.success_rate <= self.SWEET_SPOT_MAX and t.total_attempts > 0] untested = [t for t in all_states if t.total_attempts == 0] mastered = [t for t in all_states if t.status == "mastered"] if self.current_iteration <= 3: weights = {"sweet": 0.0, "explore": 1.0, "retention": 0.0} elif self.current_iteration <= 10: weights = {"sweet": 0.50, "explore": 0.35, "retention": 0.15} else: weights = {"sweet": 0.60, "explore": 0.25, "retention": 0.15} if self._detect_plateau(): weights["explore"] = min(0.50, weights["explore"] + 0.2) weights["sweet"] = max(0.2, weights["sweet"] - 0.2) # Start with minimum allocation for ALL topics (5% split) MIN_ALLOCATION = 0.05 probs: Dict[str, float] = {t.topic_name: MIN_ALLOCATION / len(all_states) for t in all_states} remaining_mass = 1.0 - MIN_ALLOCATION bonus_probs: Dict[str, float] = {} # Allocate sweet spot budget with within-iteration diversity penalty if sweet_spot: # Apply diversity penalty based on current iteration attempts staleness = {} for t in sweet_spot: # Strong penalty for topics sampled many times in current iteration if t.current_iteration_attempts == 0: # Not yet sampled this iteration - highest priority staleness[t.topic_name] = 10.0 elif t.current_iteration_attempts <= 3: # Sampled 1-3 times - moderate priority staleness[t.topic_name] = 5.0 / t.current_iteration_attempts else: # Sampled 4+ times - heavily penalized staleness[t.topic_name] = 1.0 / (t.current_iteration_attempts ** 1.5) total_stale = sum(staleness.values()) if total_stale > 0: for t in sweet_spot: bonus_probs[t.topic_name] = bonus_probs.get(t.topic_name, 0.0) + ( remaining_mass * weights["sweet"] * (staleness[t.topic_name] / total_stale) ) # Allocate explore budget - ensure we always explore something explore_pool = untested if untested else self._get_diverse_exploration_pool(sweet_spot) if explore_pool: each = remaining_mass * weights["explore"] / len(explore_pool) for t in explore_pool: bonus_probs[t.topic_name] = bonus_probs.get(t.topic_name, 0.0) + each # Allocate retention budget retention_due = [t for t in mastered if self._schedule_retention_test(t) <= self.current_iteration] if retention_due: each = remaining_mass * weights["retention"] / len(retention_due) for t in retention_due: bonus_probs[t.topic_name] = bonus_probs.get(t.topic_name, 0.0) + each # Add bonus to base minimum allocation for topic, bonus in bonus_probs.items(): probs[topic] = probs.get(topic, 0.0) + bonus # Normalize to ensure sum = 1.0 total = sum(probs.values()) if total <= 0: each = 1.0 / len(all_states) return {t.topic_name: each for t in all_states} normalized = {topic: value / total for topic, value in probs.items()} # Apply topic probability floor to prevent mode collapse MIN_TOPIC_PROB = 0.02 # Every topic gets at least 2% chance for topic in normalized: if normalized[topic] < MIN_TOPIC_PROB: normalized[topic] = MIN_TOPIC_PROB # Re-normalize after applying floor total = sum(normalized.values()) normalized = {topic: value / total for topic, value in normalized.items()} # Log top 5 topics for debugging top_topics = sorted(normalized.items(), key=lambda x: x[1], reverse=True)[:5] logger.debug(f"Topic probabilities: {top_topics}") return normalized def _get_boundary_topics(self) -> List[TopicState]: result = [] for state in self.topics.values(): if state.total_attempts == 0: continue near_low = abs(state.success_rate - self.SWEET_SPOT_MIN) <= 0.08 near_high = abs(state.success_rate - self.SWEET_SPOT_MAX) <= 0.08 if near_low or near_high: result.append(state) if not result: result = sorted(self.topics.values(), key=lambda t: abs(t.success_rate - self.TARGET_SUCCESS))[:4] return result def _get_diverse_exploration_pool(self, exclude_sweet_spot: List[TopicState]) -> List[TopicState]: """ Get topics for exploration that are NOT in sweet spot. Prioritizes: 1. Under-practiced topics (low attempt count) 2. Topics with potential (success rate 0.2-0.4 or 0.7-0.9) 3. Topics not recently attempted Args: exclude_sweet_spot: Topics already in sweet spot to exclude Returns: List of 3-5 topics for exploration """ sweet_spot_names = {t.topic_name for t in exclude_sweet_spot} candidates = [t for t in self.topics.values() if t.topic_name not in sweet_spot_names] if not candidates: # Fallback if somehow all topics are in sweet spot return list(self.topics.values())[:3] # Score each candidate scored = [] for state in candidates: # Factor 1: Under-practiced (inverse of attempts) attempt_score = 1.0 / (1.0 + state.total_attempts / 10.0) # Factor 2: Near sweet spot boundaries (could improve into sweet spot) if 0.2 <= state.success_rate < self.SWEET_SPOT_MIN: potential_score = 2.0 # Just below sweet spot - high potential elif self.SWEET_SPOT_MAX < state.success_rate <= 0.9: potential_score = 1.5 # Just above sweet spot - could be challenged more elif state.total_attempts == 0: potential_score = 3.0 # Untested - highest priority else: potential_score = 0.5 # Far from sweet spot # Factor 3: Staleness (not practiced recently) staleness_score = max(1, self.current_iteration - state.last_practiced) / 5.0 # Combined score total_score = attempt_score + potential_score + staleness_score scored.append((state, total_score)) # Return top 3-5 topics by score scored.sort(key=lambda x: x[1], reverse=True) num_explore = min(5, max(3, len(candidates) // 3)) return [state for state, _ in scored[:num_explore]] def _get_difficulty_for_topic(self, topic: str) -> float: state = self.topics[topic] noise = random.uniform(-0.04, 0.04) return max(0.1, min(0.95, state.difficulty_target + noise)) def _difficulty_to_step_range(self, difficulty: float) -> str: if difficulty < 0.3: return "1-2" if difficulty < 0.6: return "2-3" return "3-4" def _schedule_retention_test(self, state: TopicState) -> int: if state.mastered_at_iteration is None: return 10 ** 9 interval = min(2 ** max(0, state.retention_tests_passed), 32) return state.mastered_at_iteration + interval def _run_retention_tests_if_due(self) -> None: for state in self.topics.values(): if state.status != "mastered": continue if self._schedule_retention_test(state) <= self.current_iteration: # Retention test scheduling is represented by increasing pressure # during topic selection rather than explicit immediate update. logger.info("Topic %s is due for retention test", state.topic_name) def handle_retention_test_result(self, topic: str, success_rate: float) -> None: state = self.topics[topic] if success_rate >= 0.7: state.retention_tests_passed += 1 state.last_retention_score = success_rate state.status = "mastered" elif success_rate >= 0.4: state.retention_tests_passed = 0 state.last_retention_score = success_rate state.status = "active" else: state.retention_tests_passed = 0 state.last_retention_score = success_rate state.status = "forgotten" state.difficulty_target = max(0.15, state.difficulty_target * 0.7) def handle_persistent_failure(self, topic: str) -> None: state = self.topics[topic] failures = state.consecutive_failures if failures >= 3: state.difficulty_target = max(0.1, state.difficulty_target * 0.6) if failures >= 5: state.status = "paused" if failures >= 10: hard_topics = [ t for t in self.topics.values() if t.total_attempts >= 10 and t.success_rate < 0.3 ] if len(hard_topics) >= 3: self._emergency_reset() def _emergency_reset(self) -> None: logger.warning("Emergency curriculum reset triggered") for state in self.topics.values(): state.status = "active" state.difficulty_target = min(0.45, max(0.2, state.difficulty_target)) state.consecutive_failures = 0 def _detect_plateau(self) -> bool: if len(self.recent_combined_rewards) < 10: return False window = self.recent_combined_rewards[-10:] return float(np.std(window)) < 0.05