Update env.py

env.py

@@ -1,457 +1,381 @@
-"""Core RL logic for clinical trial patient screening."""
+"""Core RL environment for clinical trial patient screening."""
 
 from __future__ import annotations
 
-from copy import deepcopy
-from dataclasses import dataclass, field
-from typing import Any, Dict, List, Sequence, Set
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional
 from uuid import uuid4
 
-from openenv.core.env_server.types import State
-
-from models import (
-    ClinicalTrialScreeningAction,
-    ClinicalTrialScreeningObservation,
-    RewardModel,
-    TaskDifficulty,
-)
+from openenv.core.env_server.interfaces import Environment
+from pydantic import BaseModel, Field
+
+try:
+    from .models import (
+        ClinicalTrialAction,
+        ClinicalTrialObservation,
+        ClinicalTrialReward,
+        ClinicalTrialState,
+    )
+except ImportError:
+    from models import (
+        ClinicalTrialAction,
+        ClinicalTrialObservation,
+        ClinicalTrialReward,
+        ClinicalTrialState,
+    )
 
 INCREMENTAL_REWARD = 0.20
 FINAL_REWARD = 1.00
 HALLUCINATION_PENALTY = -0.50
+TASK_SEQUENCE = ["easy", "medium", "hard"]
+MIN_STRICT_SCORE = 0.01
+MAX_STRICT_SCORE = 0.99
+DEFAULT_GRADER_SCORE = 0.5
+
+
+def _normalize(value: Optional[str]) -> str:
+    return " ".join((value or "").strip().lower().replace("_", " ").split())
+
+
+class GroundTruth(BaseModel):
+    """Deterministic grader targets for a scenario."""
 
+    extracted_fields: Dict[str, str] = Field(default_factory=dict)
+    ranking: List[str] = Field(default_factory=list)
+    final_decision: str
+
+
+class ScenarioSpec(BaseModel):
+    """Scenario loaded from patient_data.json."""
 
-@dataclass(frozen=True)
-class ScreeningTask:
     task_id: str
-    difficulty: TaskDifficulty
+    difficulty: str
     title: str
-    brief: str
-    prompt: str
-    extraction_targets: Dict[str, str]
-    expected_decision: str
-    trial_metadata: Dict[str, Any]
-    ranking_ground_truth: List[str] = field(default_factory=list)
-    exclusion_ground_truth: Set[str] = field(default_factory=set)
-
-
-@dataclass
-class TaskRunState:
-    extracted_points: Dict[str, str] = field(default_factory=dict)
-    granted_rewards: Set[str] = field(default_factory=set)
-    final_submitted: bool = False
-    latest_grader_score: float = 0.0
-
-
-def _build_tasks() -> List[ScreeningTask]:
-    medium_scores = {
-        "P-M101": 0.88,
-        "P-M102": 0.71,
-        "P-M103": 0.54,
-    }
-    return [
-        ScreeningTask(
-            task_id="easy_eligibility",
-            difficulty=TaskDifficulty.EASY,
-            title="Phase II EGFR-Mutated NSCLC Eligibility Check",
-            brief="Determine whether the candidate meets five binary enrollment criteria.",
-            prompt=(
-                "Trial CT-NSCLC-201 enrolls adults with metastatic EGFR exon 19 or L858R "
-                "non-small cell lung cancer after first-line osimertinib. Candidate E-001 is "
-                "47 years old with biopsy-proven metastatic lung adenocarcinoma, EGFR exon 19 "
-                "deletion, ECOG 1, no active brain metastases, and adequate hepatic function. "
-                "Binary criteria: age >=18, confirmed metastatic NSCLC, sensitizing EGFR "
-                "mutation present, ECOG 0-1, no active CNS disease."
-            ),
-            extraction_targets={
-                "age": "47",
-                "diagnosis": "metastatic nsclc",
-                "biomarker": "egfr exon 19 deletion",
-                "ecog": "1",
-                "active_cns_disease": "no",
-            },
-            expected_decision="eligible",
-            trial_metadata={
-                "trial_id": "CT-NSCLC-201",
-                "specialty": "thoracic oncology",
-                "binary_criteria": [
-                    "adult patient",
-                    "metastatic NSCLC confirmed",
-                    "sensitizing EGFR mutation",
-                    "ECOG 0-1",
-                    "no active CNS disease",
-                ],
-            },
-        ),
-        ScreeningTask(
-            task_id="medium_patient_ranking",
-            difficulty=TaskDifficulty.MEDIUM,
-            title="Rank Patients for TROP2 ADC Expansion Cohort",
-            brief="Rank three real-world candidates by protocol fit-score for an EGFR-mutated NSCLC study.",
-            prompt=(
-                "Trial CT-LUNG-312 is an antibody-drug conjugate study for metastatic EGFR-mutated "
-                "NSCLC after progression on osimertinib. Rank candidates by expected screening fit. "
-                "P-M101: 56 years, EGFR exon 19 deletion, post-osimertinib only, ECOG 0, stable "
-                "treated brain metastases, CrCl 82 mL/min, AST/ALT normal. P-M102: 63 years, EGFR "
-                "L858R, post-osimertinib and platinum, ECOG 1, mild AST elevation 1.4x ULN, no "
-                "brain metastases, CrCl 68. P-M103: 59 years, exon 20 insertion, ECOG 1, chronic "
-                "prednisone 15 mg, recent palliative radiation 5 days ago, CrCl 61. Internal fit "
-                f"scores are predetermined as {medium_scores}."
-            ),
-            extraction_targets={
-                "P-M101_fit_score": "0.88",
-                "P-M102_fit_score": "0.71",
-                "P-M103_fit_score": "0.54",
-                "best_candidate": "P-M101",
-                "lowest_candidate": "P-M103",
-            },
-            expected_decision="P-M101>P-M102>P-M103",
-            ranking_ground_truth=["P-M101", "P-M102", "P-M103"],
-            trial_metadata={
-                "trial_id": "CT-LUNG-312",
-                "specialty": "thoracic oncology",
-                "ranking_rule": "higher fit-score ranks earlier",
-                "fit_scores": medium_scores,
-            },
-        ),
-        ScreeningTask(
-            task_id="hard_protocol_deviations",
-            difficulty=TaskDifficulty.HARD,
-            title="Identify Protocol Deviations from Unstructured Screening Note",
-            brief="Extract exclusions and protocol deviations from a realistic unstructured chart note.",
-            prompt=(
-                "Trial CT-LYMPH-440 is a CD19 bispecific study for relapsed diffuse large B-cell "
-                "lymphoma. Exclusions include prednisone >10 mg/day within 7 days, live vaccine "
-                "within 30 days, active hepatitis B viremia, ANC <1.0 x10^9/L, and major surgery "
-                "within 14 days. Screening note: 'Mr. R is a 68-year-old man with relapsed DLBCL. "
-                "He received a shingles live-attenuated vaccine 12 days ago at his PCP visit. He "
-                "remains on prednisone 20 mg daily for COPD flare and underwent laparoscopic "
-                "cholecystectomy 9 days ago. Labs today: ANC 0.9, HBV DNA undetectable on entecavir, "
-                "bilirubin normal. Team asks whether any items trigger screen failure or protocol "
-                "deviation before scheduling first dose.'"
-            ),
-            extraction_targets={
-                "age": "68",
-                "live_vaccine_days": "12",
-                "prednisone_mg": "20",
-                "surgery_days": "9",
-                "anc": "0.9",
-            },
-            expected_decision="exclude",
-            exclusion_ground_truth={
-                "live_vaccine_within_30_days",
-                "prednisone_over_10mg",
-                "major_surgery_within_14_days",
-                "anc_below_1.0",
-            },
-            trial_metadata={
-                "trial_id": "CT-LYMPH-440",
-                "specialty": "hematologic malignancy",
-                "expected_exclusion_schema": [
-                    "live_vaccine_within_30_days",
-                    "prednisone_over_10mg",
-                    "major_surgery_within_14_days",
-                    "anc_below_1.0",
-                    "active_hbv_viremia",
-                ],
-            },
-        ),
-    ]
-
-
-def _normalize(value: str | None) -> str:
-    return "" if value is None else value.strip().lower()
-
-
-class EasyEligibilityGrader:
-    @staticmethod
-    def grade(task: ScreeningTask, task_state: TaskRunState, decision: str | None = None) -> float:
-        extracted = sum(
-            1
-            for field_name, expected in task.extraction_targets.items()
-            if _normalize(task_state.extracted_points.get(field_name)) == _normalize(expected)
-        )
-        extraction_score = extracted / len(task.extraction_targets)
-        decision_score = 1.0 if _normalize(decision) == _normalize(task.expected_decision) else 0.0
-        return round((0.5 * extraction_score) + (0.5 * decision_score), 4)
-
-
-class MediumRankingGrader:
-    @staticmethod
-    def grade(task: ScreeningTask, task_state: TaskRunState, ranking: Sequence[str] | None = None) -> float:
-        extracted = sum(
-            1
-            for field_name, expected in task.extraction_targets.items()
-            if _normalize(task_state.extracted_points.get(field_name)) == _normalize(expected)
-        )
-        extraction_score = extracted / len(task.extraction_targets)
-        ranking = list(ranking or [])
-        if len(ranking) != len(task.ranking_ground_truth):
-            ranking_score = 0.0
-        else:
-            correct_positions = sum(
-                1
-                for observed, expected in zip(ranking, task.ranking_ground_truth)
-                if observed == expected
-            )
-            ranking_score = correct_positions / len(task.ranking_ground_truth)
-        return round((0.4 * extraction_score) + (0.6 * ranking_score), 4)
-
-
-class HardDeviationGrader:
-    @staticmethod
-    def grade(
-        task: ScreeningTask,
-        task_state: TaskRunState,
-        exclusions: Sequence[str] | None = None,
-        decision: str | None = None,
-    ) -> float:
-        extracted = sum(
-            1
-            for field_name, expected in task.extraction_targets.items()
-            if _normalize(task_state.extracted_points.get(field_name)) == _normalize(expected)
-        )
-        extraction_score = extracted / len(task.extraction_targets)
-        predicted = {_normalize(item) for item in exclusions or [] if item}
-        if task.exclusion_ground_truth:
-            exclusion_score = len(predicted & task.exclusion_ground_truth) / len(task.exclusion_ground_truth)
-        else:
-            exclusion_score = 0.0
-        decision_score = 1.0 if _normalize(decision) == _normalize(task.expected_decision) else 0.0
-        return round((0.3 * extraction_score) + (0.4 * exclusion_score) + (0.3 * decision_score), 4)
+    instructions: str
+    context: Dict[str, Any]
+    ground_truth: GroundTruth
+    hidden_exclusions: List[str] = Field(default_factory=list)
+    max_steps: int = 6
+    grader_name: str = "deterministic_json_grader"
+
 
+class ClinicalTrialEnvironment(
+    Environment[ClinicalTrialAction, ClinicalTrialObservation, ClinicalTrialState]
+):
+    """Clinical trial screening environment backed by externalized JSON scenarios."""
 
-class ClinicalTrialScreeningEnv:
-    """Stateful environment covering easy, medium, and hard screening tasks."""
+    SUPPORTS_CONCURRENT_SESSIONS = True
 
     def __init__(self) -> None:
-        self._tasks = _build_tasks()
-        self._state = State(episode_id=str(uuid4()), step_count=0)
-        self._index = 0
-        self._task_runs: Dict[str, TaskRunState] = {}
-        self._episode_complete = False
-
-    def reset(self) -> ClinicalTrialScreeningObservation:
-        self._state = State(episode_id=str(uuid4()), step_count=0)
-        self._index = 0
-        self._episode_complete = False
-        self._task_runs = {task.task_id: TaskRunState() for task in self._tasks}
+        super().__init__()
+        data_path = Path(__file__).resolve().with_name("patient_data.json")
+        payload = json.loads(data_path.read_text(encoding="utf-8"))
+        task_payload = payload.get("tasks", {})
+        self._scenarios: Dict[str, ScenarioSpec] = {
+            task_id: ScenarioSpec.model_validate({**scenario, "task_id": task_id})
+            for task_id, scenario in task_payload.items()
+        }
+        self._task_cursor = -1
+        self._current_scenario: Optional[ScenarioSpec] = None
+        self._submitted_ranking: List[str] = []
+        self._state = ClinicalTrialState(episode_id=str(uuid4()), step_count=0)
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        task_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> ClinicalTrialObservation:
+        del seed, kwargs
+        selected_task_id = task_id or self._next_task_id()
+        self._current_scenario = self._scenarios[selected_task_id]
+        self._submitted_ranking = []
+        self._state = ClinicalTrialState(
+            episode_id=episode_id or str(uuid4()),
+            step_count=0,
+            current_task_id=self._current_scenario.task_id,
+            difficulty=self._current_scenario.difficulty,
+            title=self._current_scenario.title,
+            extracted_fields={},
+            identified_deviations=[],
+            final_decision=None,
+            grading_score=DEFAULT_GRADER_SCORE,
+        )
         return self._build_observation(
-            reward_model=RewardModel(total=0.0),
-            feedback="Episode reset. Start with the easy eligibility assessment.",
+            reward_details=ClinicalTrialReward(
+                notes=["Episode reset."], grader_score=DEFAULT_GRADER_SCORE
+            ),
+            done=False,
         )
 
-    def step(self, action: ClinicalTrialScreeningAction) -> ClinicalTrialScreeningObservation:
-        self._state.step_count += 1
-        task = self._current_task()
-        task_state = self._task_runs[task.task_id]
-        reward = RewardModel()
-        feedback_parts: List[str] = []
-
-        if self._episode_complete:
-            reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("episode_already_complete")
-            reward.total = reward.penalty
-            return self._build_observation(reward, "Episode already completed. Reset to start a new run.")
+    def step(
+        self,
+        action: ClinicalTrialAction,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> ClinicalTrialObservation:
+        del timeout_s, kwargs
+        if self._current_scenario is None:
+            return self.reset()
 
-        if action.action_type == "destructive_action":
+        self._state.step_count += 1
+        reward = ClinicalTrialReward()
+        done = False
+        terminal_reason: Optional[str] = None
+
+        if action.action_type == "extract_data":
+            self._handle_extraction(action, reward)
+        elif action.action_type == "flag_deviation":
+            self._handle_deviation_flag(action, reward)
+        elif action.action_type == "rank_patients":
+            self._handle_ranking(action, reward)
+            if action.ranking:
+                done = True
+                terminal_reason = "ranking_submitted"
+        elif action.action_type == "submit_decision":
+            self._state.final_decision = action.final_decision
+            done = True
+            terminal_reason = "final_decision_submitted"
+        elif action.action_type == "delete_evidence":
             reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("destructive_action")
-            feedback_parts.append("Destructive action blocked in screening workflow.")
-        elif action.action_type == "extract_data":
-            feedback_parts.append(self._handle_extract_data(task, task_state, action, reward))
-        elif action.action_type == "submit_ranking":
-            feedback_parts.append(self._handle_submit_ranking(task, task_state, action, reward))
-        elif action.action_type == "flag_exclusions":
-            feedback_parts.append(self._handle_flag_exclusions(task, task_state, action, reward))
-        elif action.action_type == "final_decision":
-            feedback_parts.append(self._handle_final_decision(task, task_state, action, reward))
+            reward.notes.append("Destructive action: deleting evidence is not allowed.")
         else:
             reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("unsupported_action")
-            feedback_parts.append("Unsupported action type for this environment.")
+            reward.notes.append(f"Unsupported action type: {action.action_type}")
+
+        if self._state.step_count >= self._current_scenario.max_steps and not done:
+            done = True
+            terminal_reason = "max_steps_reached"
 
-        reward.total = round(
-            reward.incremental_reward + reward.final_reward + reward.penalty,
-            4,
+        reward.grader_score = self._grade_for_current_task()
+        self._state.grading_score = reward.grader_score
+
+        if done:
+            if self._is_final_submission_correct():
+                reward.final_reward = FINAL_REWARD
+                reward.notes.append("Correct final screening decision.")
+            reward.missing_items = self._missing_items()
+
+        reward.total_reward = round(
+            reward.incremental_reward + reward.final_reward + reward.penalty, 4
+        )
+        return self._build_observation(
+            reward_details=reward,
+            done=done,
+            terminal_reason=terminal_reason,
         )
-        return self._build_observation(reward, " ".join(part for part in feedback_parts if part))
 
-    def state(self) -> State:
+    @property
+    def state(self) -> ClinicalTrialState:
         return self._state
 
-    def _current_task(self) -> ScreeningTask:
-        return self._tasks[self._index]
+    def grader(self) -> float:
+        """Deterministically compare agent outputs against the current scenario ground truth."""
+        if self._current_scenario is None:
+            return DEFAULT_GRADER_SCORE
+        components: List[float] = []
+        truth = self._current_scenario.ground_truth
 
-    def _build_observation(
-        self,
-        reward_model: RewardModel,
-        feedback: str,
-    ) -> ClinicalTrialScreeningObservation:
-        task = self._current_task()
-        task_state = self._task_runs.get(task.task_id, TaskRunState())
-        missing_targets = [
-            field_name
-            for field_name in task.extraction_targets
-            if field_name not in task_state.granted_rewards
-        ]
-        return ClinicalTrialScreeningObservation(
-            task_id=task.task_id,
-            difficulty=task.difficulty,
-            title=task.title,
-            brief=task.brief,
-            prompt=task.prompt,
-            extracted_points=deepcopy(task_state.extracted_points),
-            missing_targets=missing_targets,
-            available_actions=[
-                "extract_data",
-                "submit_ranking" if task.difficulty is TaskDifficulty.MEDIUM else "flag_exclusions",
-                "final_decision",
-                "destructive_action",
-            ],
-            grader_score=task_state.latest_grader_score,
-            reward_breakdown=reward_model,
-            feedback=feedback,
-            done=self._episode_complete,
-            reward=reward_model.total,
-            trial_metadata=deepcopy(task.trial_metadata),
-        )
+        if truth.extracted_fields:
+            field_hits = sum(
+                1
+                for field_name, expected in truth.extracted_fields.items()
+                if _normalize(self._state.extracted_fields.get(field_name)) == _normalize(expected)
+            )
+            score = field_hits / len(truth.extracted_fields)
+            # Clamp component to ensure it never hits exact 0.0 or 1.0
+            score = min(max(score, MIN_STRICT_SCORE), MAX_STRICT_SCORE)
+            components.append(score)
 
-    def _handle_extract_data(
-        self,
-        task: ScreeningTask,
-        task_state: TaskRunState,
-        action: ClinicalTrialScreeningAction,
-        reward: RewardModel,
-    ) -> str:
-        field_name = action.field_name or ""
-        if field_name not in task.extraction_targets:
-            reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("hallucinated_field")
-            return f"Field '{field_name}' is not part of the protocol data model."
-
-        expected = task.extraction_targets[field_name]
-        observed = _normalize(action.value)
-        if observed == _normalize(expected):
-            task_state.extracted_points[field_name] = action.value or ""
-            if field_name not in task_state.granted_rewards:
-                task_state.granted_rewards.add(field_name)
-                reward.incremental_reward += INCREMENTAL_REWARD
-                reward.reasons.append(f"correct_extraction:{field_name}")
-            task_state.latest_grader_score = self._current_grader_score(task, task_state, action)
-            return f"Captured {field_name} correctly."
+        if self._current_scenario.hidden_exclusions:
+            exclusion_hits = sum(
+                1
+                for exclusion in self._current_scenario.hidden_exclusions
+                if exclusion in self._state.identified_deviations
+            )
+            score = exclusion_hits / len(self._current_scenario.hidden_exclusions)
+            # Clamp component to ensure it never hits exact 0.0 or 1.0
+            score = min(max(score, MIN_STRICT_SCORE), MAX_STRICT_SCORE)
+            components.append(score)
+
+        if truth.ranking:
+            ranking = self._submitted_ranking
+            if ranking and len(ranking) == len(truth.ranking):
+                positional_hits = sum(
+                    1 for actual, expected in zip(ranking, truth.ranking) if actual == expected
+                ) / len(truth.ranking)
+                pairwise_hits = 0
+                total_pairs = 0
+                for index, higher in enumerate(truth.ranking):
+                    for lower in truth.ranking[index + 1 :]:
+                        total_pairs += 1
+                        if ranking.index(higher) < ranking.index(lower):
+                            pairwise_hits += 1
+                pairwise_score = pairwise_hits / max(total_pairs, 1)
+                score = (0.6 * positional_hits) + (0.4 * pairwise_score)
+            else:
+                score = MIN_STRICT_SCORE  # Penalize missing/incorrect ranking
+            # Clamp component to ensure it never hits exact 0.0 or 1.0
+            score = min(max(score, MIN_STRICT_SCORE), MAX_STRICT_SCORE)
+            components.append(score)
+
+        # Final decision correctness
+        final_match = _normalize(self._state.final_decision) == _normalize(truth.final_decision)
+        score = MAX_STRICT_SCORE if final_match else MIN_STRICT_SCORE  # Already clamped
+        components.append(score)
+
+        if not components:
+            return MIN_STRICT_SCORE
+
+        raw_score = sum(components) / len(components)
+        strict_score = min(max(raw_score, MIN_STRICT_SCORE), MAX_STRICT_SCORE)
+        return round(strict_score, 4)
+
+    def grade_easy_screening(self) -> float:
+        """Task-specific grader for the easy screening task."""
+        if self._current_scenario is None or self._current_scenario.task_id != "easy":
+            self.reset(task_id="easy")
+        return self.grader()
+
+    def grade_medium_ranking(self) -> float:
+        """Task-specific grader for the medium ranking task."""
+        if self._current_scenario is None or self._current_scenario.task_id != "medium":
+            self.reset(task_id="medium")
+        return self.grader()
+
+    def grade_hard_exclusions(self) -> float:
+        """Task-specific grader for the hard exclusions task."""
+        if self._current_scenario is None or self._current_scenario.task_id != "hard":
+            self.reset(task_id="hard")
+        return self.grader()
+
+    def _grade_for_current_task(self) -> float:
+        """Resolve and run the grader declared by the current scenario."""
+        assert self._current_scenario is not None
+        grader_name = (self._current_scenario.grader_name or "").strip()
+        grader_fn = getattr(self, grader_name, None)
+        if callable(grader_fn):
+            score = float(grader_fn())
+        else:
+            score = float(self.grader())
+        return round(min(max(score, MIN_STRICT_SCORE), MAX_STRICT_SCORE), 4)
 
-        reward.penalty += HALLUCINATION_PENALTY
-        reward.reasons.append(f"incorrect_extraction:{field_name}")
-        return f"Extracted value for {field_name} does not match the chart."
+    def _next_task_id(self) -> str:
+        self._task_cursor = (self._task_cursor + 1) % len(TASK_SEQUENCE)
+        return TASK_SEQUENCE[self._task_cursor]
 
-    def _handle_submit_ranking(
-        self,
-        task: ScreeningTask,
-        task_state: TaskRunState,
-        action: ClinicalTrialScreeningAction,
-        reward: RewardModel,
-    ) -> str:
-        if task.difficulty is not TaskDifficulty.MEDIUM:
+    def _handle_extraction(self, action: ClinicalTrialAction, reward: ClinicalTrialReward) -> None:
+        assert self._current_scenario is not None
+        if not action.field_name or action.value is None:
             reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("ranking_on_non_medium_task")
-            return "Ranking is only valid on the medium task."
-
-        task_state.final_submitted = True
-        is_correct = list(action.ranking) == task.ranking_ground_truth
-        if is_correct:
-            reward.final_reward += FINAL_REWARD
-            reward.reasons.append("correct_ranking")
-        else:
+            reward.notes.append("extract_data requires field_name and value.")
+            return
+
+        expected_value = self._current_scenario.ground_truth.extracted_fields.get(action.field_name)
+        if expected_value is None:
             reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("incorrect_ranking")
-        task_state.latest_grader_score = MediumRankingGrader.grade(task, task_state, action.ranking)
-        self._advance_task()
-        return "Ranking accepted." if is_correct else "Ranking accepted but does not match the deterministic fit ordering."
+            reward.notes.append(f"Hallucinated field: {action.field_name}")
+            return
 
-    def _handle_flag_exclusions(
-        self,
-        task: ScreeningTask,
-        task_state: TaskRunState,
-        action: ClinicalTrialScreeningAction,
-        reward: RewardModel,
-    ) -> str:
-        if task.difficulty is not TaskDifficulty.HARD:
+        if _normalize(action.value) == _normalize(expected_value):
+            if action.field_name not in self._state.extracted_fields:
+                reward.incremental_reward += INCREMENTAL_REWARD
+                reward.matched_items.append(action.field_name)
+                reward.notes.append(f"Validated extraction for {action.field_name}.")
+            self._state.extracted_fields[action.field_name] = action.value
+        else:
             reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("exclusion_flag_on_non_hard_task")
-            return "Exclusion flagging is reserved for the hard chart-review task."
+            reward.notes.append(f"Incorrect value for {action.field_name}.")
 
-        predicted = {_normalize(item) for item in action.exclusions if item}
-        invalid = predicted - task.exclusion_ground_truth
-        if invalid:
+    def _handle_deviation_flag(self, action: ClinicalTrialAction, reward: ClinicalTrialReward) -> None:
+        assert self._current_scenario is not None
+        submitted = [_normalize(item) for item in action.deviations]
+        if not submitted:
             reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("hallucinated_exclusion")
-            task_state.latest_grader_score = HardDeviationGrader.grade(
-                task,
-                task_state,
-                exclusions=action.exclusions,
-            )
-            return f"Unsupported exclusion codes submitted: {sorted(invalid)}."
+            reward.notes.append("flag_deviation requires at least one deviation.")
+            return
 
-        task_state.latest_grader_score = HardDeviationGrader.grade(
-            task,
-            task_state,
-            exclusions=action.exclusions,
+        for deviation in submitted:
+            if deviation in self._current_scenario.hidden_exclusions:
+                if deviation not in self._state.identified_deviations:
+                    self._state.identified_deviations.append(deviation)
+                    reward.incremental_reward += INCREMENTAL_REWARD
+                    reward.matched_items.append(deviation)
+                    reward.notes.append(f"Validated deviation: {deviation}.")
+            else:
+                reward.penalty += HALLUCINATION_PENALTY
+                reward.notes.append(f"Unsupported deviation claim: {deviation}.")
+
+    def _handle_ranking(self, action: ClinicalTrialAction, reward: ClinicalTrialReward) -> None:
+        assert self._current_scenario is not None
+        ranking = action.ranking
+        valid_patients = [
+            patient["patient_id"]
+            for patient in self._current_scenario.context.get("patients", [])
+        ]
+        if sorted(ranking) != sorted(valid_patients):
+            reward.penalty += HALLUCINATION_PENALTY
+            reward.notes.append("Ranking must include each patient exactly once.")
+            return
+        self._submitted_ranking = ranking
+        self._state.final_decision = "ranking_submitted"
+
+    def _is_final_submission_correct(self) -> bool:
+        assert self._current_scenario is not None
+        truth = self._current_scenario.ground_truth
+        if truth.ranking:
+            return self._submitted_ranking == truth.ranking
+        return _normalize(self._state.final_decision) == _normalize(truth.final_decision)
+
+    def _missing_items(self) -> List[str]:
+        assert self._current_scenario is not None
+        truth = self._current_scenario.ground_truth
+        missing_fields = [
+            field_name
+            for field_name, expected in truth.extracted_fields.items()
+            if _normalize(self._state.extracted_fields.get(field_name)) != _normalize(expected)
+        ]
+        missing_fields.extend(
+            exclusion
+            for exclusion in self._current_scenario.hidden_exclusions
+            if exclusion not in self._state.identified_deviations
         )
-        return "Exclusion codes recorded."
+        if truth.ranking and self._submitted_ranking != truth.ranking:
+            missing_fields.append("ranking")
+        if _normalize(self._state.final_decision) != _normalize(truth.final_decision):
+            missing_fields.append("final_decision")
+        return missing_fields
 
-    def _handle_final_decision(
+    def _build_observation(
         self,
-        task: ScreeningTask,
-        task_state: TaskRunState,
-        action: ClinicalTrialScreeningAction,
-        reward: RewardModel,
-    ) -> str:
-        decision = action.value or ""
-        task_state.final_submitted = True
-        if task.difficulty is TaskDifficulty.MEDIUM:
-            reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("decision_instead_of_ranking")
-            task_state.latest_grader_score = MediumRankingGrader.grade(task, task_state)
-            return "Medium task requires submit_ranking rather than final_decision."
-
-        is_correct = _normalize(decision) == _normalize(task.expected_decision)
-        if is_correct:
-            reward.final_reward += FINAL_REWARD
-            reward.reasons.append("correct_final_decision")
-        else:
-            reward.penalty += HALLUCINATION_PENALTY
-            reward.reasons.append("incorrect_final_decision")
-
-        if task.difficulty is TaskDifficulty.EASY:
-            task_state.latest_grader_score = EasyEligibilityGrader.grade(task, task_state, decision=decision)
-        else:
-            exclusions = sorted(task.exclusion_ground_truth)
-            task_state.latest_grader_score = HardDeviationGrader.grade(
-                task,
-                task_state,
-                exclusions=exclusions,
-                decision=decision,
-            )
+        reward_details: ClinicalTrialReward,
+        done: bool,
+        terminal_reason: Optional[str] = None,
+    ) -> ClinicalTrialObservation:
+        assert self._current_scenario is not None
+        attempts_remaining = max(self._current_scenario.max_steps - self._state.step_count, 0)
+        return ClinicalTrialObservation(
+            task_id=self._current_scenario.task_id,
+            difficulty=self._current_scenario.difficulty,  # type: ignore[arg-type]
+            title=self._current_scenario.title,
+            instructions=self._current_scenario.instructions,
+            context=self._current_scenario.context,
+            expected_fields=list(self._current_scenario.ground_truth.extracted_fields.keys()),
+            extracted_fields=dict(self._state.extracted_fields),
+            identified_deviations=list(self._state.identified_deviations),
+            attempts_remaining=attempts_remaining,
+            grader_name=self._current_scenario.grader_name,
+            reward_details=reward_details,
+            reward=reward_details.total_reward,
+            done=done,
+            metadata={"grading_score": self._state.grading_score},
+            terminal_reason=terminal_reason,
+        )
 
-        self._advance_task()
-        return "Final screening decision accepted." if is_correct else "Final decision conflicts with protocol evidence."
 
-    def _current_grader_score(
-        self,
-        task: ScreeningTask,
-        task_state: TaskRunState,
-        action: ClinicalTrialScreeningAction,
-    ) -> float:
-        if task.difficulty is TaskDifficulty.EASY:
-            return EasyEligibilityGrader.grade(task, task_state)
-        if task.difficulty is TaskDifficulty.MEDIUM:
-            return MediumRankingGrader.grade(task, task_state)
-        return HardDeviationGrader.grade(task, task_state, exclusions=action.exclusions)
-
-    def _advance_task(self) -> None:
-        if self._index == len(self._tasks) - 1:
-            self._episode_complete = True
-            return
-        self._index += 1
+class ClinicalTrialEnv(ClinicalTrialEnvironment):
+    """Compatibility alias for manifest entry points expecting env:ClinicalTrialEnv."""
 
+    pass