""" server/environment.py — Core medical diagnostic environment logic. This is the BRAIN of the Medical Diagnostic Environment. It: 1. Manages patient cases and episode state 2. Processes agent actions (questions, tests, diagnoses) 3. Calculates rewards based on diagnostic quality 4. Provides trajectory-based reward signals (not sparse) Pure Python - no HTTP or WebSocket code here. All logic is deterministic and reproducible. """ import random import uuid from typing import Dict, List, Optional from openenv.core.env_server import Environment import sys import os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from models import DiagnosticAction, PatientObservation, ClinicalState from server.medical_data import ( PATIENT_CASES, calculate_question_reward, calculate_test_reward, calculate_diagnosis_accuracy, get_patient_response, normalize_test_name, ) class MedicalDiagnosticEnvironment(Environment): """ Medical Diagnostic Environment for RL Training. Simulates doctor-patient interaction where an LLM agent must: 1. Ask relevant clinical questions 2. Order appropriate diagnostic tests 3. Make accurate diagnoses The environment provides rich reward signals throughout the trajectory: - +0.05 per relevant question asked - +0.10 per informative test ordered - +1.0 for correct final diagnosis - Penalizes inefficient or irrelevant actions This is NOT a sparse reward environment - the agent sees meaningful progress at each step, which is crucial for learning. """ SUPPORTS_CONCURRENT_SESSIONS = True # Allow multiple parallel training sessions # Episode configuration MAX_STEPS = 15 # Maximum actions per episode DIFFICULTY_LEVELS = ["easy", "medium", "hard"] def __init__(self): """Initialize environment state.""" super().__init__() # Episode state variables self._episode_id: str = "" self._case_id: str = "" self._difficulty: str = "" self._step_count: int = 0 self._total_reward: float = 0.0 # Patient interaction tracking self._questions_asked: List[str] = [] self._tests_ordered: List[str] = [] self._test_results: Dict = {} self._diagnosis_submitted: Optional[str] = None self._final_accuracy: float = 0.0 # Episode status self._done: bool = False self._episode_reward_breakdown: Dict = { "question_rewards": 0.0, "test_rewards": 0.0, "diagnosis_reward": 0.0, "efficiency_penalty": 0.0, } # Initialize a default episode so stateless HTTP /step can operate # on a fresh environment when the server creates a new instance. self.reset(difficulty="easy") @property def current_case_id(self) -> str: """Current patient case identifier.""" return self._case_id @property def current_difficulty(self) -> str: """Current episode difficulty level.""" return self._difficulty # ───────────────────────────────────────────────────────────────────── # Core API Methods # ───────────────────────────────────────────────────────────────────── def reset(self, difficulty: str = "easy", **kwargs) -> PatientObservation: """ Reset the environment for a new diagnostic episode. Args: difficulty: "easy", "medium", or "hard" (controls case selection) Returns: Initial PatientObservation for the agent to read """ # Initialize episode self._episode_id = str(uuid.uuid4()) self._difficulty = difficulty if difficulty in self.DIFFICULTY_LEVELS else "easy" self._case_id = self._select_case_by_difficulty(self._difficulty) self._step_count = 0 self._total_reward = 0.0 # Reset tracking self._questions_asked = [] self._tests_ordered = [] self._test_results = {} self._diagnosis_submitted = None self._final_accuracy = 0.0 self._done = False self._episode_reward_breakdown = { "question_rewards": 0.0, "test_rewards": 0.0, "diagnosis_reward": 0.0, "efficiency_penalty": 0.0, } # Get case information case = PATIENT_CASES[self._case_id] # Create initial observation initial_message = ( f"Patient presents with: {case['presentation']}\n" f"Age: {case['age']}, Gender: {case['gender']}\n" f"You have up to {self.MAX_STEPS} steps to diagnose this patient.\n" f"Please start by asking questions or ordering tests." ) return PatientObservation( done=False, reward=0.0, message=initial_message, patient_response=None, test_result=None, questions_asked=[], tests_completed=[], patient_data_revealed={ "age": case["age"], "gender": case["gender"], "presentation": case["presentation"], }, steps_taken=0, max_steps=self.MAX_STEPS, ) def step(self, action: DiagnosticAction, **kwargs) -> PatientObservation: """ Process one diagnostic action (question, test, or diagnosis). Returns immediate reward and next observation. """ if self._done: return self._create_done_observation( message="Episode already ended. Call reset() to start a new case." ) self._step_count += 1 step_reward = 0.0 message = "" patient_response = None test_result = None # ── Process action based on type ── if action.action_type == "ask_question": step_reward, message, patient_response = self._handle_question(action.question) elif action.action_type == "order_test": step_reward, message, test_result = self._handle_test(action.test_name) elif action.action_type == "submit_diagnosis": step_reward, message = self._handle_diagnosis(action.diagnosis) self._done = True else: message = f"Unknown action type: {action.action_type}" step_reward = -0.05 # Accumulate rewards self._total_reward += step_reward # Check if episode should end if self._step_count >= self.MAX_STEPS and not self._done: message += f"\nMax steps reached. Episode ending." self._done = True # Get current case for patient data revelation case = PATIENT_CASES[self._case_id] return PatientObservation( done=self._done, reward=step_reward, message=message, patient_response=patient_response, test_result=test_result, questions_asked=self._questions_asked.copy(), tests_completed=self._tests_ordered.copy(), patient_data_revealed=self._build_patient_data_view(case), steps_taken=self._step_count, max_steps=self.MAX_STEPS, ) @property def state(self) -> ClinicalState: """ Return complete internal state (includes hidden information). Used for debugging only - NEVER send to agent. """ case = PATIENT_CASES.get(self._case_id, {}) return ClinicalState( episode_id=self._episode_id, step_count=self._step_count, true_diagnosis=case.get("true_diagnosis", ""), patient_case=self._case_id, patient_id=self._case_id, patient_details=case, difficulty=self._difficulty, questions_asked=self._questions_asked.copy(), tests_completed=self._tests_ordered.copy(), final_diagnosis_submitted=self._diagnosis_submitted, final_accuracy=self._final_accuracy, ) # ───────────────────────────────────────────────────────────────────── # Action Processing # ───────────────────────────────────────────────────────────────────── def _handle_question(self, question: Optional[str]) -> tuple: """ Process a question about the patient. Returns: (reward, message, patient_response) """ if not question or not isinstance(question, str) or not question.strip(): message = "No valid question was provided. Please ask a clinical question." return -0.05, message, None # Calculate reward for asking this question reward = calculate_question_reward(self._case_id, question) # Record question self._questions_asked.append(question) self._episode_reward_breakdown["question_rewards"] += reward # Get patient response response = get_patient_response(self._case_id, question) message = f"Patient response: {response}" if reward == 0.00: message += " (Question may not be directly relevant)" elif reward == 0.01: message += " (Somewhat relevant question)" else: message += " (Good clinical question!)" return reward, message, response def _handle_test(self, test_name: Optional[str]) -> tuple: """ Process a test order. Returns: (reward, message, test_result_dict) """ if not test_name or not isinstance(test_name, str) or not test_name.strip(): message = "No valid test name was provided. Please order a valid diagnostic test." return -0.05, message, None # Calculate reward for ordering this test reward = calculate_test_reward(self._case_id, test_name) # Get case data case = PATIENT_CASES[self._case_id] # Try to find matching test result test_result_data = None matched_test_key = None test_lower = normalize_test_name(test_name) for test_key, result in case.get("test_results", {}).items(): if test_key.lower() == test_lower or test_key.lower() in test_lower or test_lower in test_key.lower(): test_result_data = result matched_test_key = test_key break if test_result_data is None: message = f"Test '{test_name}' not available for this patient or unavailable in this setting." reward = -0.02 return reward, message, None # Record test self._tests_ordered.append(matched_test_key) self._test_results[matched_test_key] = test_result_data self._episode_reward_breakdown["test_rewards"] += reward # Format test result for agent test_result_dict = { "test_name": matched_test_key, "result": str(test_result_data), "interpretation": test_result_data.get("interpretation", test_result_data.get("finding", "")) } message = f"Test result received for {matched_test_key}:\n{test_result_dict['interpretation']}" if reward == 0.10: message += " (Excellent diagnostic test!)" elif reward == 0.05: message += " (Useful supporting test)" else: message += " (Test ordered but may be less relevant)" return reward, message, test_result_dict def _handle_diagnosis(self, diagnosis: str) -> tuple: """ Process final diagnosis submission. Returns: (reward, message) """ # Calculate diagnostic accuracy accuracy = calculate_diagnosis_accuracy(self._case_id, diagnosis) self._final_accuracy = accuracy self._diagnosis_submitted = diagnosis # Create diagnosis reward (not just accuracy, but also process quality) case = PATIENT_CASES[self._case_id] true_diagnosis = case["true_diagnosis"] # Change the if/elif chain to use >= comparisons: if accuracy >= 0.95: reward = 1.0 message = f"Correct diagnosis: {diagnosis}" elif accuracy >= 0.7: reward = accuracy message = f"Acceptable diagnosis: {diagnosis}. True: {true_diagnosis}" elif accuracy >= 0.3: reward = accuracy message = f"Partially correct. True: {true_diagnosis}" else: reward = 0.0 message = f"Incorrect. True: {true_diagnosis}" self._episode_reward_breakdown["diagnosis_reward"] = reward # Add efficiency feedback if self._step_count > self.MAX_STEPS * 0.8: penalty = 0.1 * (self._step_count / self.MAX_STEPS - 0.8) self._episode_reward_breakdown["efficiency_penalty"] = penalty message += f"\n(Efficiency penalty: -{penalty:.2f} for taking many steps)" return reward, message # ───────────────────────────────────────────────────────────────────── # Helper Methods # ───────────────────────────────────────────────────────────────────── def _select_case_by_difficulty(self, difficulty: str) -> str: matching_keys = [k for k, v in PATIENT_CASES.items() if v["difficulty"] == difficulty] if not matching_keys: matching_keys = list(PATIENT_CASES.keys()) return random.choice(matching_keys) def _build_patient_data_view(self, case: Dict) -> Dict: """ Build what the agent has learned about the patient so far. Only includes information revealed through questions/tests. """ revealed = { "age": case.get("age"), "gender": case.get("gender"), "presentation": case.get("presentation"), } # Add findings based on questions asked findings = case.get("hidden_findings", {}) for question in self._questions_asked: q_lower = question.lower() for finding, value in findings.items(): if finding.lower().replace("_", " ") in q_lower: revealed[f"finding_{finding}"] = value # Add test results if self._test_results: revealed["test_results"] = self._test_results return revealed def _create_done_observation(self, message: str) -> PatientObservation: """Create a terminal observation.""" return PatientObservation( done=True, reward=0.0, message=message, patient_response=None, test_result=None, questions_asked=self._questions_asked.copy(), tests_completed=self._tests_ordered.copy(), patient_data_revealed={}, steps_taken=self._step_count, max_steps=self.MAX_STEPS, ) def get_episode_summary(self) -> Dict: """ Return a summary of the episode for logging/evaluation. """ case = PATIENT_CASES.get(self._case_id, {}) return { "episode_id": self._episode_id, "case_id": self._case_id, "difficulty": self._difficulty, "true_diagnosis": case.get("true_diagnosis", ""), "submitted_diagnosis": self._diagnosis_submitted, "accuracy": self._final_accuracy, "diagnostic_accuracy": self._final_accuracy, "total_reward": self._total_reward, "steps": self._step_count, "steps_taken": self._step_count, "max_steps": self.MAX_STEPS, "questions_asked": len(self._questions_asked), "tests_ordered": len(self._tests_ordered), "reward_breakdown": self._episode_reward_breakdown, }