Upload 3 files

README.md

@@ -1,135 +1,9 @@
----
-title: Clinical Trial Patient Screening
-emoji: 🧪
-colorFrom: green
-colorTo: indigo
-sdk: docker
-pinned: false
-app_port: 8000
-tags:
-  - openenv
-short_description: RL environment for clinical trial patient screening.
----
-# Clinical Trial Patient Screening Environment
+# Clinical Trial Screening OpenEnv
 
-OpenEnv environment for clinical trial patient screening with three deterministic healthcare tasks:
+This project implements a realistic reinforcement-learning environment for clinical trial patient screening with three tasks:
 
-- `easy`: binary eligibility screening against 5 criteria
-- `medium`: ranking 3 patients by protocol fit
-- `hard`: protocol deviation and exclusion detection from unstructured clinical text
+- `easy`: eligibility determination against five binary protocol criteria
+- `medium`: ranking three candidates for an EGFR-mutated NSCLC trial
+- `hard`: detecting exclusions and protocol deviations from unstructured chart text
 
-This environment is designed as a real-world screening workflow rather than a toy game. It uses typed Pydantic models, deterministic programmatic graders, incremental reward shaping for correct data extraction, and terminal rewards for correct screening outcomes.
-
-## Task Overview
-
-### Easy
-EGFR-mutated metastatic NSCLC eligibility check using structured oncology data.
-
-### Medium
-Rank 3 HER2-positive metastatic breast cancer candidates by fit for a trial.
-
-### Hard
-Detect subtle exclusions from an AML screening note, including protocol deviations such as recent investigational treatment, active infection, QTc prolongation, and CYP3A4 inhibitor exposure.
-
-## Reward Design
-
-- `+0.20` for each correct clinical data point or valid deviation extracted
-- `+1.00` for a correct final screening decision
-- `-0.50` for hallucinated fields, invalid deviation claims, or destructive actions
-
-Each task also produces a deterministic grader score in `[0.0, 1.0]`.
-
-## Project Structure
-
-```text
-clinical_trial_env/
-├── .env
-├── Dockerfile
-├── README.md
-├── __init__.py
-├── client.py
-├── env.py
-├── inference.py
-├── models.py
-├── openenv.yaml
-├── pyproject.toml
-└── server/
-    ├── __init__.py
-    ├── app.py
-    └── clinical_trial_env_environment.py
-```
-
-## Build And Run
-
-Build the container from the project root:
-
-```bash
-docker build -t clinical-trial-env:latest .
-```
-
-Run the server locally:
-
-```bash
-docker run --rm -p 8000:8000 clinical-trial-env:latest
-```
-
-Validate the environment:
-
-```bash
-openenv validate .
-```
-
-## Inference
-
-The root [inference.py](/Users/abhishekkanade/Desktop/Hackathon/OpenEnv/clinical_trial_env/inference.py) uses the OpenAI client and emits exactly:
-
-- `[START]`
-- `[STEP]`
-- `[END]`
-
-Required environment variables:
-
-- `HF_TOKEN`
-- `LOCAL_IMAGE_NAME` or `ENV_BASE_URL`
-- `API_BASE_URL` optional, defaults to Hugging Face router
-- `MODEL_NAME` optional
-- `CLINICAL_TRIAL_TASK` with values `easy`, `medium`, or `hard`
-
-Example:
-
-```bash
-set -a
-source .env
-set +a
-python3 inference.py
-```
-
-## Python Usage
-
-```python
-import asyncio
-
-from clinical_trial_env import ClinicalTrialAction, ClinicalTrialEnv
-
-
-async def main() -> None:
-    env = await ClinicalTrialEnv.from_docker_image("clinical-trial-env:latest")
-    try:
-        result = await env.reset(task_id="easy")
-        result = await env.step(
-            ClinicalTrialAction(action_type="extract_data", field_name="age", value="56")
-        )
-        print(result.reward, result.observation.reward_details.grader_score)
-    finally:
-        await env.close()
-
-
-asyncio.run(main())
-```
-
-## Audit Notes
-
-- `reset()`, `step(action)`, and OpenEnv `state` access are implemented in [env.py](/Users/abhishekkanade/Desktop/Hackathon/OpenEnv/clinical_trial_env/env.py).
-- Gold answers are not exposed through observation metadata.
-- Graders are deterministic and bounded in `[0.0, 1.0]`.
-- The hard task uses unstructured medical text and clinically realistic exclusion criteria.
+Use `openenv validate` to validate the environment and `python inference.py` for a local scripted run.

env.py

@@ -1,348 +1,457 @@
-"""Core RL environment for clinical trial patient screening."""
+"""Core RL logic for clinical trial patient screening."""
 
 from __future__ import annotations
 
-import json
-from pathlib import Path
-from typing import Any, Dict, List, Optional
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Sequence, Set
 from uuid import uuid4
 
-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,
-    )
+from openenv.core.env_server.types import State
+
+from models import (
+    ClinicalTrialScreeningAction,
+    ClinicalTrialScreeningObservation,
+    RewardModel,
+    TaskDifficulty,
+)
 
 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
-
-
-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: str
+    difficulty: TaskDifficulty
     title: str
-    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"
-
+    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)
 
-class ClinicalTrialEnvironment(
-    Environment[ClinicalTrialAction, ClinicalTrialObservation, ClinicalTrialState]
-):
-    """Clinical trial screening environment backed by externalized JSON scenarios."""
 
-    SUPPORTS_CONCURRENT_SESSIONS = True
+class ClinicalTrialScreeningEnv:
+    """Stateful environment covering easy, medium, and hard screening tasks."""
 
     def __init__(self) -> None:
-        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=0.5,
-        )
+        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}
         return self._build_observation(
-            reward_details=ClinicalTrialReward(notes=["Episode reset."], grader_score=0.5),
-            done=False,
+            reward_model=RewardModel(total=0.0),
+            feedback="Episode reset. Start with the easy eligibility assessment.",
         )
 
-    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()
-
+    def step(self, action: ClinicalTrialScreeningAction) -> ClinicalTrialScreeningObservation:
         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":
+        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.notes.append("Destructive action: deleting evidence is not allowed.")
+            reward.reasons.append("episode_already_complete")
+            reward.total = reward.penalty
+            return self._build_observation(reward, "Episode already completed. Reset to start a new run.")
+
+        if action.action_type == "destructive_action":
+            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))
         else:
             reward.penalty += HALLUCINATION_PENALTY
-            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.grader_score = self.grader()
-        self._state.grading_score = reward.grader_score
+            reward.reasons.append("unsupported_action")
+            feedback_parts.append("Unsupported action type for this environment.")
 
-        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,
+        reward.total = round(
+            reward.incremental_reward + reward.final_reward + reward.penalty,
+            4,
         )
+        return self._build_observation(reward, " ".join(part for part in feedback_parts if part))
 
-    @property
-    def state(self) -> ClinicalTrialState:
+    def state(self) -> State:
         return self._state
 
-    def grader(self) -> float:
-        """Deterministically compare agent outputs against the current scenario ground truth."""
-        assert self._current_scenario is not None
-        components: List[float] = []
-        truth = self._current_scenario.ground_truth
+    def _current_task(self) -> ScreeningTask:
+        return self._tasks[self._index]
 
-        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 _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 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 _next_task_id(self) -> str:
-        self._task_cursor = (self._task_cursor + 1) % len(TASK_SEQUENCE)
-        return TASK_SEQUENCE[self._task_cursor]
-
-    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:
+    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.notes.append("extract_data requires field_name and value.")
-            return
+            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."
 
-        expected_value = self._current_scenario.ground_truth.extracted_fields.get(action.field_name)
-        if expected_value is None:
-            reward.penalty += HALLUCINATION_PENALTY
-            reward.notes.append(f"Hallucinated field: {action.field_name}")
-            return
+        reward.penalty += HALLUCINATION_PENALTY
+        reward.reasons.append(f"incorrect_extraction:{field_name}")
+        return f"Extracted value for {field_name} does not match the chart."
 
-        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
+    def _handle_submit_ranking(
+        self,
+        task: ScreeningTask,
+        task_state: TaskRunState,
+        action: ClinicalTrialScreeningAction,
+        reward: RewardModel,
+    ) -> str:
+        if task.difficulty is not TaskDifficulty.MEDIUM:
+            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.penalty += HALLUCINATION_PENALTY
-            reward.notes.append(f"Incorrect value for {action.field_name}.")
+            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."
 
-    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:
+    def _handle_flag_exclusions(
+        self,
+        task: ScreeningTask,
+        task_state: TaskRunState,
+        action: ClinicalTrialScreeningAction,
+        reward: RewardModel,
+    ) -> str:
+        if task.difficulty is not TaskDifficulty.HARD:
             reward.penalty += HALLUCINATION_PENALTY
-            reward.notes.append("flag_deviation requires at least one deviation.")
-            return
+            reward.reasons.append("exclusion_flag_on_non_hard_task")
+            return "Exclusion flagging is reserved for the hard chart-review task."
 
-        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):
+        predicted = {_normalize(item) for item in action.exclusions if item}
+        invalid = predicted - task.exclusion_ground_truth
+        if invalid:
             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
+            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)}."
+
+        task_state.latest_grader_score = HardDeviationGrader.grade(
+            task,
+            task_state,
+            exclusions=action.exclusions,
         )
-        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
+        return "Exclusion codes recorded."
 
-    def _build_observation(
+    def _handle_final_decision(
         self,
-        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,
-        )
+        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,
+            )
 
-class ClinicalTrialEnv(ClinicalTrialEnvironment):
-    """Compatibility alias for manifest entry points expecting env:ClinicalTrialEnv."""
+        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
 
-    pass

inference.py

@@ -1,43 +1,32 @@
-"""Hackathon-compliant inference runner for the clinical trial environment."""
+"""Benchmark-compatible inference runner for clinical trial screening."""
 
 from __future__ import annotations
 
 import asyncio
 import json
 import os
-import textwrap
-from typing import Dict, List, Optional, Tuple
+import re
+from typing import List, Optional
 
 from openai import OpenAI
 
-try:
-    from clinical_trial_env import ClinicalTrialAction, ClinicalTrialEnv
-except ImportError:
-    from client import ClinicalTrialEnv
-    from models import ClinicalTrialAction
+from client import ClinicalTrialScreeningEnvClient
+from models import ClinicalTrialScreeningAction, ClinicalTrialScreeningObservation
 
-LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") or os.getenv("IMAGE_NAME")
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-TASK_NAME = os.getenv("CLINICAL_TRIAL_TASK", "easy")
-BENCHMARK = os.getenv("CLINICAL_TRIAL_BENCHMARK", "clinical_trial_env")
+HF_TOKEN = os.getenv("HF_TOKEN")
+LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
 ENV_BASE_URL = os.getenv("ENV_BASE_URL")
-MAX_STEPS = int(os.getenv("MAX_STEPS", "20"))
-TEMPERATURE = float(os.getenv("TEMPERATURE", "0.1"))
-MAX_TOKENS = int(os.getenv("MAX_TOKENS", "220"))
-SUCCESS_SCORE_THRESHOLD = float(os.getenv("SUCCESS_SCORE_THRESHOLD", "0.8"))
-
-SYSTEM_PROMPT = textwrap.dedent(
-    """
-    You are operating a clinical trial screening environment.
-    Return exactly one compact JSON object with keys:
-    action_type, field_name, value, ranking, deviations, final_decision, rationale.
-    Use only supported action_type values:
-    extract_data, rank_patients, flag_deviation, submit_decision.
-    Do not add markdown, commentary, or code fences.
-    """
-).strip()
+BENCHMARK = os.getenv("BENCHMARK", "clinical_trial_screening")
+TASK_NAME = os.getenv("TASK_NAME", "clinical_trial_patient_screening")
+MAX_STEPS = int(os.getenv("MAX_STEPS", "19"))
+
+SYSTEM_PROMPT = (
+    "You are a clinical trial screening agent. "
+    "Return one compact JSON object with keys action_type, target_id, field_name, value, "
+    "ranking, exclusions, rationale. Use only protocol-supported fields and codes."
+)
 
 
 def log_start(task: str, env: str, model: str) -> None:
@@ -45,191 +34,201 @@ def log_start(task: str, env: str, model: str) -> None:
 
 
 def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
-    error_text = error if error is not None else "null"
+    error_value = error if error is not None else "null"
     print(
         f"[STEP] step={step} action={action} reward={reward:.2f} "
-        f"done={str(done).lower()} error={error_text}",
+        f"done={str(done).lower()} error={error_value}",
         flush=True,
     )
 
 
-def log_end(success: bool, steps: int, rewards: List[float]) -> None:
+def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
     rewards_str = ",".join(f"{reward:.2f}" for reward in rewards)
     print(
-        f"[END] success={str(success).lower()} steps={steps} rewards={rewards_str}",
+        f"[END] success={str(success).lower()} steps={steps} "
+        f"score={score:.2f} rewards={rewards_str}",
         flush=True,
     )
 
 
-def sanitize_error(error: Optional[str]) -> Optional[str]:
-    if error is None:
+def format_action(action: ClinicalTrialScreeningAction) -> str:
+    if action.action_type == "extract_data":
+        return f"extract_data({action.field_name}={action.value})"
+    if action.action_type == "submit_ranking":
+        return f"submit_ranking({'>'.join(action.ranking)})"
+    if action.action_type == "flag_exclusions":
+        return f"flag_exclusions({','.join(action.exclusions)})"
+    if action.action_type == "final_decision":
+        return f"final_decision({action.value})"
+    return action.action_type
+
+
+def safe_error(message: Optional[str]) -> Optional[str]:
+    if not message:
         return None
-    cleaned = " ".join(error.split())
-    return cleaned or "null"
-
-
-def build_user_prompt(task_name: str, step: int, observation_payload: Dict, history: List[str]) -> str:
-    history_text = "\n".join(history[-4:]) if history else "None"
-    return textwrap.dedent(
-        f"""
-        Task: {task_name}
-        Step: {step}
-        Observation:
-        {json.dumps(observation_payload, indent=2, sort_keys=True)}
-
-        Recent history:
-        {history_text}
-
-        Return the next best JSON action.
-        """
-    ).strip()
-
-
-def heuristic_action(task_name: str, step: int) -> ClinicalTrialAction:
-    heuristics: Dict[Tuple[str, int], ClinicalTrialAction] = {
-        ("easy", 1): ClinicalTrialAction(action_type="extract_data", field_name="age", value="56"),
-        ("easy", 2): ClinicalTrialAction(
-            action_type="extract_data", field_name="egfr_mutation", value="L858R positive"
-        ),
-        ("easy", 3): ClinicalTrialAction(action_type="submit_decision", final_decision="eligible"),
-        ("medium", 1): ClinicalTrialAction(
-            action_type="extract_data", field_name="BC-101_her2_status", value="IHC 3+"
-        ),
-        ("medium", 2): ClinicalTrialAction(
-            action_type="extract_data", field_name="BC-102_trastuzumab_exposure", value="none"
-        ),
-        ("medium", 3): ClinicalTrialAction(
-            action_type="rank_patients", ranking=["BC-101", "BC-103", "BC-102"]
-        ),
-        ("hard", 1): ClinicalTrialAction(action_type="extract_data", field_name="biomarker", value="FLT3-ITD"),
-        ("hard", 2): ClinicalTrialAction(
-            action_type="flag_deviation",
-            deviations=[
-                "neutropenic fever",
-                "qtc greater than 480 ms",
-                "recent strong CYP3A4 inhibitor",
-            ],
-        ),
-        ("hard", 3): ClinicalTrialAction(action_type="submit_decision", final_decision="ineligible"),
-    }
-    return heuristics.get((task_name, step), ClinicalTrialAction(action_type="submit_decision", final_decision="ineligible"))
+    return re.sub(r"\s+", " ", message.strip())
+
+
+def build_user_prompt(
+    observation: ClinicalTrialScreeningObservation,
+    history: List[str],
+) -> str:
+    recent_history = " | ".join(history[-4:]) if history else "none"
+    return (
+        f"task_id={observation.task_id}\n"
+        f"difficulty={observation.difficulty.value}\n"
+        f"title={observation.title}\n"
+        f"brief={observation.brief}\n"
+        f"prompt={observation.prompt}\n"
+        f"missing_targets={observation.missing_targets}\n"
+        f"available_actions={observation.available_actions}\n"
+        f"trial_metadata={json.dumps(observation.trial_metadata, sort_keys=True)}\n"
+        f"history={recent_history}\n"
+        "Reply with JSON only."
+    )
 
 
-def parse_action(raw_text: str) -> ClinicalTrialAction:
-    payload = json.loads(raw_text)
-    return ClinicalTrialAction.model_validate(payload)
+def fallback_action(observation: ClinicalTrialScreeningObservation) -> ClinicalTrialScreeningAction:
+    task_id = observation.task_id
+    extracted = observation.extracted_points
+    if task_id == "easy_eligibility":
+        for field_name, value in [
+            ("age", "47"),
+            ("diagnosis", "metastatic nsclc"),
+            ("biomarker", "egfr exon 19 deletion"),
+            ("ecog", "1"),
+            ("active_cns_disease", "no"),
+        ]:
+            if field_name not in extracted:
+                return ClinicalTrialScreeningAction(
+                    action_type="extract_data",
+                    field_name=field_name,
+                    value=value,
+                )
+        return ClinicalTrialScreeningAction(action_type="final_decision", value="eligible")
+    if task_id == "medium_patient_ranking":
+        for field_name, value in [
+            ("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"),
+        ]:
+            if field_name not in extracted:
+                return ClinicalTrialScreeningAction(
+                    action_type="extract_data",
+                    field_name=field_name,
+                    value=value,
+                )
+        return ClinicalTrialScreeningAction(
+            action_type="submit_ranking",
+            ranking=["P-M101", "P-M102", "P-M103"],
+        )
+    for field_name, value in [
+        ("age", "68"),
+        ("live_vaccine_days", "12"),
+        ("prednisone_mg", "20"),
+        ("surgery_days", "9"),
+        ("anc", "0.9"),
+    ]:
+        if field_name not in extracted:
+            return ClinicalTrialScreeningAction(
+                action_type="extract_data",
+                field_name=field_name,
+                value=value,
+            )
+    if observation.grader_score < 0.7:
+        return ClinicalTrialScreeningAction(
+            action_type="flag_exclusions",
+            exclusions=[
+                "live_vaccine_within_30_days",
+                "prednisone_over_10mg",
+                "major_surgery_within_14_days",
+                "anc_below_1.0",
+            ],
+        )
+    return ClinicalTrialScreeningAction(action_type="final_decision", value="exclude")
 
 
 def get_model_action(
     client: OpenAI,
-    task_name: str,
-    step: int,
-    observation_payload: Dict,
+    observation: ClinicalTrialScreeningObservation,
     history: List[str],
-) -> ClinicalTrialAction:
-    user_prompt = build_user_prompt(task_name, step, observation_payload, history)
+) -> tuple[ClinicalTrialScreeningAction, Optional[str]]:
+    prompt = build_user_prompt(observation, history)
     try:
-        completion = client.chat.completions.create(
+        response = client.chat.completions.create(
             model=MODEL_NAME,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
+                {"role": "user", "content": prompt},
             ],
-            temperature=TEMPERATURE,
-            max_tokens=MAX_TOKENS,
-            stream=False,
+            temperature=0.0,
+            max_tokens=200,
         )
-        content = (completion.choices[0].message.content or "").strip()
-        return parse_action(content)
-    except Exception:
-        return heuristic_action(task_name, step)
+        content = (response.choices[0].message.content or "").strip()
+        payload = json.loads(content)
+        return ClinicalTrialScreeningAction.model_validate(payload), None
+    except Exception as exc:
+        return fallback_action(observation), safe_error(str(exc))
 
 
-def format_action(action: ClinicalTrialAction) -> str:
-    payload = {
-        "action_type": action.action_type,
-        "field_name": action.field_name,
-        "value": action.value,
-        "ranking": action.ranking,
-        "deviations": action.deviations,
-        "final_decision": action.final_decision,
-    }
-    return json.dumps(payload, separators=(",", ":"), sort_keys=True)
-
-
-async def create_env() -> ClinicalTrialEnv:
+async def create_env_client() -> ClinicalTrialScreeningEnvClient:
     if LOCAL_IMAGE_NAME:
-        return await ClinicalTrialEnv.from_docker_image(LOCAL_IMAGE_NAME)
+        return await ClinicalTrialScreeningEnvClient.from_docker_image(LOCAL_IMAGE_NAME)
     if ENV_BASE_URL:
-        env = ClinicalTrialEnv(base_url=ENV_BASE_URL)
-        await env.connect()
-        return env
-    raise RuntimeError("Set LOCAL_IMAGE_NAME for Docker execution or ENV_BASE_URL for an existing server.")
+        client = ClinicalTrialScreeningEnvClient(base_url=ENV_BASE_URL)
+        return await client.connect()
+    client = ClinicalTrialScreeningEnvClient(base_url="http://localhost:8000")
+    return await client.connect()
 
 
 async def main() -> None:
-    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-    env: Optional[ClinicalTrialEnv] = None
+    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+    env = await create_env_client()
     rewards: List[float] = []
+    history: List[str] = []
     steps_taken = 0
-    score = 0.0
+    final_score = 0.0
     success = False
-    history: List[str] = []
     last_error: Optional[str] = None
+    result = None
 
     log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
 
     try:
-        env = await create_env()
-        result = await env.reset(task_id=TASK_NAME)
-
+        result = await env.reset()
         for step in range(1, MAX_STEPS + 1):
             if result.done:
                 break
 
-            action = get_model_action(
-                client=client,
-                task_name=TASK_NAME,
-                step=step,
-                observation_payload=result.observation.model_dump(mode="json"),
-                history=history,
-            )
-
-            try:
-                result = await env.step(action)
-                reward = float(result.reward or 0.0)
-                done = bool(result.done)
-                last_error = None
-            except Exception as exc:
-                reward = 0.0
-                done = True
-                last_error = sanitize_error(str(exc))
+            action, planning_error = get_model_action(client, result.observation, history)
+            result = await env.step(action)
 
+            reward = float(result.reward or 0.0)
             rewards.append(reward)
             steps_taken = step
+            last_error = planning_error
             log_step(
                 step=step,
                 action=format_action(action),
                 reward=reward,
-                done=done,
-                error=sanitize_error(last_error),
+                done=bool(result.done),
+                error=last_error,
             )
-            history.append(f"step={step} action={format_action(action)} reward={reward:.2f}")
-
-            if last_error is not None or done:
+            history.append(
+                f"{result.observation.task_id}:{format_action(action)}:{reward:.2f}:{result.done}"
+            )
+            if result.done:
                 break
 
-        if last_error is None and "result" in locals():
-            score = float(result.observation.reward_details.grader_score)
-        score = max(0.0, min(score, 1.0))
-        success = last_error is None and score >= SUCCESS_SCORE_THRESHOLD
+        if result is not None:
+            final_score = float(result.observation.grader_score)
+            final_score = min(max(final_score, 0.0), 1.0)
+            success = bool(result.done) and final_score >= 0.99
     finally:
-        if env is not None:
-            try:
-                await env.close()
-            except Exception as exc:
-                last_error = last_error or sanitize_error(str(exc))
-        log_end(success=success, steps=steps_taken, rewards=rewards)
+        await env.close()
+        log_end(success=success, steps=steps_taken, score=final_score, rewards=rewards)
 
 
 if __name__ == "__main__":