shivamtech9395/maths_reasoning_env
0
1"""2MathReasoningEnvironment — core server-side logic.3Implements reset(), step(), state() with 5 graded task types.4"""5 6import uuid7import random8import math9from typing import Tuple10 11from ..models import MathAction, MathObservation, MathState12 13 14# ──────────────────────────────────────────────15# Graders (one per task type — at least 3 required)16# ──────────────────────────────────────────────17 18def grade_arithmetic(problem_data: dict, answer: str) -> Tuple[bool, str]:19 """Grader for arithmetic tasks."""20 try:21 user_val = float(answer.strip().replace(",", ""))22 expected = float(problem_data["answer"])23 correct = abs(user_val - expected) < 1e-624 return correct, str(problem_data["answer"])25 except (ValueError, KeyError):26 return False, str(problem_data.get("answer", ""))27 28 29def grade_algebra(problem_data: dict, answer: str) -> Tuple[bool, str]:30 """Grader for algebra tasks (linear equations)."""31 try:32 user_val = float(answer.strip())33 expected = float(problem_data["answer"])34 correct = abs(user_val - expected) < 1e-435 return correct, str(problem_data["answer"])36 except (ValueError, KeyError):37 return False, str(problem_data.get("answer", ""))38 39 40def grade_word_problems(problem_data: dict, answer: str) -> Tuple[bool, str]:41 """Grader for word problems."""42 try:43 user_val = float(answer.strip().replace(",", ""))44 expected = float(problem_data["answer"])45 correct = abs(user_val - expected) < 1e-446 return correct, str(problem_data["answer"])47 except (ValueError, KeyError):48 return False, str(problem_data.get("answer", ""))49 50 51def grade_number_theory(problem_data: dict, answer: str) -> Tuple[bool, str]:52 """Grader for number theory tasks (GCD, LCM, prime check)."""53 try:54 user_val = int(answer.strip())55 expected = int(problem_data["answer"])56 correct = user_val == expected57 return correct, str(problem_data["answer"])58 except (ValueError, KeyError):59 return False, str(problem_data.get("answer", ""))60 61 62def grade_geometry(problem_data: dict, answer: str) -> Tuple[bool, str]:63 """Grader for geometry tasks (area, perimeter, volume)."""64 try:65 user_val = float(answer.strip())66 expected = float(problem_data["answer"])67 correct = abs(user_val - expected) < 0.0168 return correct, str(round(problem_data["answer"], 4))69 except (ValueError, KeyError):70 return False, str(problem_data.get("answer", ""))71 72 73# Map task_type → grader function74GRADERS = {75 "arithmetic": grade_arithmetic,76 "algebra": grade_algebra,77 "word_problems": grade_word_problems,78 "number_theory": grade_number_theory,79 "geometry": grade_geometry,80}81 82TASK_TYPES = list(GRADERS.keys())83 84 85# ──────────────────────────────────────────────86# Problem generators87# ──────────────────────────────────────────────88 89def generate_arithmetic():90 ops = ['+', '-', '*', '/']91 op = random.choice(ops)92 a = random.randint(1, 50)93 b = random.randint(1, 20)94 if op == '+':95 ans = a + b96 elif op == '-':97 ans = a - b98 elif op == '*':99 ans = a * b100 else:101 b = random.randint(1, 10)102 ans = round(a / b, 4)103 return {"problem": f"What is {a} {op} {b}?", "answer": ans}104 105 106def generate_algebra():107 # ax + b = c → x = (c - b) / a108 a = random.randint(2, 10)109 b = random.randint(1, 20)110 x = random.randint(-10, 10)111 c = a * x + b112 return {"problem": f"Solve for x: {a}x + {b} = {c}", "answer": x}113 114 115def generate_word_problem():116 templates = [117 lambda: {118 "problem": f"Alice has {random.randint(5,30)} apples. She gives {random.randint(1,4)} to each of her {random.randint(2,5)} friends. How many apples does she have left?",119 "answer": None # computed below120 },121 ]122 n_apples = random.randint(10, 50)123 give_each = random.randint(1, 4)124 n_friends = random.randint(2, 5)125 left = n_apples - give_each * n_friends126 if left < 0:127 left = 0128 return {129 "problem": f"Alice has {n_apples} apples. She gives {give_each} to each of her {n_friends} friends. How many apples does she have left?",130 "answer": left131 }132 133 134def generate_number_theory():135 task = random.choice(["gcd", "lcm"])136 a = random.randint(2, 30)137 b = random.randint(2, 30)138 if task == "gcd":139 ans = math.gcd(a, b)140 return {"problem": f"What is the GCD of {a} and {b}?", "answer": ans}141 else:142 ans = (a * b) // math.gcd(a, b)143 return {"problem": f"What is the LCM of {a} and {b}?", "answer": ans}144 145 146def generate_geometry():147 shape = random.choice(["circle_area", "rectangle_area", "triangle_area", "rectangle_perimeter"])148 if shape == "circle_area":149 r = random.randint(1, 10)150 ans = round(math.pi * r * r, 4)151 return {"problem": f"What is the area of a circle with radius {r}? (Use π=3.14159, round to 4 decimal places)", "answer": ans}152 elif shape == "rectangle_area":153 w, h = random.randint(2, 20), random.randint(2, 20)154 return {"problem": f"What is the area of a rectangle with width {w} and height {h}?", "answer": w * h}155 elif shape == "triangle_area":156 b, h = random.randint(2, 20), random.randint(2, 20)157 ans = 0.5 * b * h158 return {"problem": f"What is the area of a triangle with base {b} and height {h}?", "answer": ans}159 else:160 w, h = random.randint(2, 20), random.randint(2, 20)161 return {"problem": f"What is the perimeter of a rectangle with width {w} and height {h}?", "answer": 2 * (w + h)}162 163 164GENERATORS = {165 "arithmetic": generate_arithmetic,166 "algebra": generate_algebra,167 "word_problems": generate_word_problem,168 "number_theory": generate_number_theory,169 "geometry": generate_geometry,170}171 172 173# ──────────────────────────────────────────────174# Environment class175# ──────────────────────────────────────────────176 177class MathReasoningEnvironment:178 """179 OpenEnv-compatible Math Reasoning Environment.180 Supports 5 graded task types. Gymnasium-style API: reset / step / state.181 """182 183 def __init__(self):184 self._state = MathState()185 self._current_problem_data: dict = {}186 187 def reset(self, task_type: str = None) -> MathObservation:188 task = task_type if task_type in TASK_TYPES else random.choice(TASK_TYPES)189 self._state = MathState(190 episode_id=str(uuid.uuid4()),191 current_task=task,192 max_steps=10,193 )194 problem_data = GENERATORS[task]()195 self._current_problem_data = problem_data196 return MathObservation(197 problem=problem_data["problem"],198 task_type=task,199 feedback=None,200 correct=None,201 correct_answer=None,202 episode_done=False,203 score=0.0,204 )205 206 def step(self, action: MathAction) -> MathObservation:207 self._state.step_count += 1208 self._state.problems_attempted += 1209 210 task = self._state.current_task211 grader = GRADERS[task]212 correct, correct_answer = grader(self._current_problem_data, action.answer)213 214 # Reward shaping215 reward = 1.0 if correct else -0.2216 # Bonus for chain-of-thought reasoning217 if action.reasoning and len(action.reasoning.strip()) > 10:218 reward += 0.1219 self._state.total_score += reward220 221 if correct:222 self._state.problems_correct += 1223 feedback = "✅ Correct! Well done."224 else:225 feedback = f"❌ Incorrect. The correct answer was {correct_answer}."226 227 done = self._state.step_count >= self._state.max_steps228 229 # Generate next problem (unless done)230 if not done:231 next_task = random.choice(TASK_TYPES)232 self._state.current_task = next_task233 next_problem = GENERATORS[next_task]()234 self._current_problem_data = next_problem235 next_problem_text = next_problem["problem"]236 else:237 next_problem_text = "Episode complete."238 239 return MathObservation(240 problem=next_problem_text,241 task_type=self._state.current_task,242 feedback=feedback,243 correct=correct,244 correct_answer=correct_answer if not correct else None,245 episode_done=done,246 score=round(self._state.total_score, 4),247 )248 249 @property250 def state(self) -> MathState:251 return self._state252