shivamtech9395/maths_reasoning_env

1"""2Pydantic models for MathReasoningEnv.3Defines Action, Observation, and State contracts.4"""5 6from typing import Optional, Literal7from pydantic import BaseModel, Field8 9 10class MathAction(BaseModel):11    """Action sent by the agent: an answer string to a math problem."""12    answer: str = Field(..., description="The agent's answer to the current math problem")13    reasoning: Optional[str] = Field(None, description="Optional chain-of-thought reasoning")14 15 16class MathObservation(BaseModel):17    """Observation returned after each step."""18    problem: str = Field(..., description="The current math problem")19    task_type: str = Field(..., description="Type of task: arithmetic/algebra/word_problems/number_theory/geometry")20    feedback: Optional[str] = Field(None, description="Feedback on the last answer (if any)")21    correct: Optional[bool] = Field(None, description="Whether the last answer was correct")22    correct_answer: Optional[str] = Field(None, description="The correct answer (shown after wrong attempt)")23    episode_done: bool = Field(False, description="Whether the episode is finished")24    score: float = Field(0.0, description="Running score in this episode")25 26 27class MathState(BaseModel):28    """Internal state of the environment."""29    episode_id: str = Field("", description="Unique episode ID")30    step_count: int = Field(0, description="Number of steps taken")31    total_score: float = Field(0.0, description="Accumulated reward score")32    current_task: str = Field("arithmetic", description="Current task type")33    max_steps: int = Field(10, description="Max steps per episode")34    problems_attempted: int = Field(0, description="Number of problems attempted")35    problems_correct: int = Field(0, description="Number of problems answered correctly")36