Upload 14 files

Dockerfile

@@ -0,0 +1,48 @@
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY . /app/env
+WORKDIR /app/env
+
+RUN if ! command -v uv >/dev/null 2>&1; then \
+        curl -LsSf https://astral.sh/uv/install.sh | sh && \
+        mv /root/.local/bin/uv /usr/local/bin/uv && \
+        mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-install-project --no-editable; \
+    else \
+        uv sync --no-install-project --no-editable; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-editable; \
+    else \
+        uv sync --no-editable; \
+    fi
+
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+COPY --from=builder /app/env/.venv /app/.venv
+COPY --from=builder /app/env /app/env
+COPY --from=builder /app/env/patient_data.json /app/patient_data.json
+
+ENV PATH="/app/.venv/bin:$PATH"
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+ENV ENABLE_WEB_INTERFACE=true
+
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

README.md

@@ -6,8 +6,130 @@ colorTo: indigo
 sdk: docker
 pinned: false
 app_port: 8000
-base_path: /web
 tags:
   - openenv
 short_description: RL environment for clinical trial patient screening.
----
+---
+# Clinical Trial Patient Screening Environment
+
+OpenEnv environment for clinical trial patient screening with three deterministic healthcare 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
+
+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.

__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Clinical Trial screening environment."""
+
+from .client import ClinicalTrialEnv
+from .env import ClinicalTrialEnvironment
+from .models import (
+    ClinicalTrialAction,
+    ClinicalTrialObservation,
+    ClinicalTrialReward,
+    ClinicalTrialState,
+)
+
+__all__ = [
+    "ClinicalTrialAction",
+    "ClinicalTrialObservation",
+    "ClinicalTrialReward",
+    "ClinicalTrialState",
+    "ClinicalTrialEnvironment",
+    "ClinicalTrialEnv",
+]

client.py

@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Clinical Trial Env environment client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+
+try:
+    from .models import ClinicalTrialAction, ClinicalTrialObservation, ClinicalTrialState
+except ImportError:
+    from models import ClinicalTrialAction, ClinicalTrialObservation, ClinicalTrialState
+
+
+class ClinicalTrialEnv(
+    EnvClient[ClinicalTrialAction, ClinicalTrialObservation, ClinicalTrialState]
+):
+    """
+    Client for the clinical trial screening environment.
+
+    This client maintains a persistent WebSocket connection to the environment server,
+    enabling efficient multi-step interactions with lower latency.
+    Each client instance has its own dedicated environment session on the server.
+
+    Example:
+        >>> # Connect to a running server
+        >>> with ClinicalTrialEnv(base_url="http://localhost:8000") as client:
+        ...     result = await client.reset(task_id="easy")
+        ...     print(result.observation.title)
+        ...
+        ...     result = await client.step(
+        ...         ClinicalTrialAction(action_type="submit_decision", final_decision="eligible")
+        ...     )
+        ...     print(result.observation.reward_details.grader_score)
+
+    Example with Docker:
+        >>> # Automatically start container and connect
+        >>> client = await ClinicalTrialEnv.from_docker_image("clinical-trial-env:latest")
+        >>> try:
+        ...     result = await client.reset(task_id="hard")
+        ...     result = await client.step(
+        ...         ClinicalTrialAction(action_type="flag_deviation", deviations=["active uncontrolled infection"])
+        ...     )
+        ... finally:
+        ...     await client.close()
+    """
+
+    def _step_payload(self, action: ClinicalTrialAction) -> Dict:
+        """
+        Convert ClinicalTrialAction to JSON payload for step message.
+
+        Args:
+            action: ClinicalTrialAction instance
+
+        Returns:
+            Dictionary representation suitable for JSON encoding
+        """
+        return action.model_dump(mode="json")
+
+    def _parse_result(self, payload: Dict) -> StepResult[ClinicalTrialObservation]:
+        """
+        Parse server response into StepResult[ClinicalTrialObservation].
+
+        Args:
+            payload: JSON response data from server
+
+        Returns:
+            StepResult with ClinicalTrialObservation
+        """
+        obs_data = payload.get("observation", {})
+        observation = ClinicalTrialObservation.model_validate(
+            {
+                **obs_data,
+                "done": obs_data.get("done", payload.get("done", False)),
+                "reward": obs_data.get("reward", payload.get("reward")),
+            }
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> ClinicalTrialState:
+        """
+        Parse server response into State object.
+
+        Args:
+            payload: JSON response from state request
+
+        Returns:
+            State object with episode_id and step_count
+        """
+        return ClinicalTrialState.model_validate(payload)

env.py

@@ -0,0 +1,328 @@
+"""Core RL environment for clinical trial patient screening."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+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,
+    )
+
+INCREMENTAL_REWARD = 0.20
+FINAL_REWARD = 1.00
+HALLUCINATION_PENALTY = -0.50
+TASK_SEQUENCE = ["easy", "medium", "hard"]
+
+
+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."""
+
+    task_id: str
+    difficulty: str
+    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"
+
+
+class ClinicalTrialEnvironment(
+    Environment[ClinicalTrialAction, ClinicalTrialObservation, ClinicalTrialState]
+):
+    """Clinical trial screening environment backed by externalized JSON scenarios."""
+
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    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.0,
+        )
+        return self._build_observation(
+            reward_details=ClinicalTrialReward(notes=["Episode reset."]),
+            done=False,
+        )
+
+    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()
+
+        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.notes.append("Destructive action: deleting evidence is not allowed.")
+        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"
+
+        if done:
+            reward.grader_score = self.grader()
+            if self._is_final_submission_correct():
+                reward.final_reward = FINAL_REWARD
+                reward.notes.append("Correct final screening decision.")
+            reward.missing_items = self._missing_items()
+            self._state.grading_score = reward.grader_score
+
+        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,
+        )
+
+    @property
+    def state(self) -> ClinicalTrialState:
+        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
+
+        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)
+            )
+            components.append(field_hits / len(truth.extracted_fields))
+
+        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
+            )
+            components.append(exclusion_hits / len(self._current_scenario.hidden_exclusions))
+
+        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)
+                components.append((0.6 * positional_hits) + (0.4 * pairwise_score))
+            else:
+                components.append(0.0)
+
+        components.append(
+            1.0
+            if _normalize(self._state.final_decision) == _normalize(truth.final_decision)
+            else 0.0
+        )
+
+        if not components:
+            return 0.0
+        return round(sum(components) / len(components), 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:
+            reward.penalty += HALLUCINATION_PENALTY
+            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.notes.append(f"Hallucinated field: {action.field_name}")
+            return
+
+        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.notes.append(f"Incorrect value for {action.field_name}.")
+
+    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.notes.append("flag_deviation requires at least one deviation.")
+            return
+
+        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
+        )
+        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 _build_observation(
+        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,
+        )

inference.py

@@ -0,0 +1,236 @@
+"""Hackathon-compliant inference runner for the clinical trial environment."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import textwrap
+from typing import Dict, List, Optional, Tuple
+
+from openai import OpenAI
+
+try:
+    from clinical_trial_env import ClinicalTrialAction, ClinicalTrialEnv
+except ImportError:
+    from client import ClinicalTrialEnv
+    from models import ClinicalTrialAction
+
+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")
+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()
+
+
+def log_start(task: str, env: str, model: str) -> None:
+    print(f"[START] task={task} env={env} model={model}", flush=True)
+
+
+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"
+    print(
+        f"[STEP] step={step} action={action} reward={reward:.2f} "
+        f"done={str(done).lower()} error={error_text}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, 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}",
+        flush=True,
+    )
+
+
+def sanitize_error(error: Optional[str]) -> Optional[str]:
+    if error is None:
+        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"))
+
+
+def parse_action(raw_text: str) -> ClinicalTrialAction:
+    payload = json.loads(raw_text)
+    return ClinicalTrialAction.model_validate(payload)
+
+
+def get_model_action(
+    client: OpenAI,
+    task_name: str,
+    step: int,
+    observation_payload: Dict,
+    history: List[str],
+) -> ClinicalTrialAction:
+    user_prompt = build_user_prompt(task_name, step, observation_payload, history)
+    try:
+        completion = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ],
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+            stream=False,
+        )
+        content = (completion.choices[0].message.content or "").strip()
+        return parse_action(content)
+    except Exception:
+        return heuristic_action(task_name, step)
+
+
+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:
+    if LOCAL_IMAGE_NAME:
+        return await ClinicalTrialEnv.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.")
+
+
+async def main() -> None:
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    env: Optional[ClinicalTrialEnv] = None
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+    history: List[str] = []
+    last_error: Optional[str] = None
+
+    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        env = await create_env()
+        result = await env.reset(task_id=TASK_NAME)
+
+        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))
+
+            rewards.append(reward)
+            steps_taken = step
+            log_step(
+                step=step,
+                action=format_action(action),
+                reward=reward,
+                done=done,
+                error=sanitize_error(last_error),
+            )
+            history.append(f"step={step} action={format_action(action)} reward={reward:.2f}")
+
+            if last_error is not None or 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
+    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)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

models.py

@@ -0,0 +1,85 @@
+"""Typed models for the clinical trial screening environment."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional
+
+from openenv.core.env_server.types import Action, Observation, State
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ClinicalTrialReward(BaseModel):
+    """Structured reward payload for deterministic grading and shaping."""
+
+    model_config = ConfigDict(
+        extra="forbid",
+        validate_assignment=True,
+        arbitrary_types_allowed=True,
+    )
+
+    incremental_reward: float = Field(default=0.0, description="Reward from validated extractions.")
+    final_reward: float = Field(default=0.0, description="Terminal reward for a correct final decision.")
+    penalty: float = Field(default=0.0, description="Penalty for hallucinations or destructive actions.")
+    total_reward: float = Field(default=0.0, description="Net reward for the step.")
+    grader_score: float = Field(default=0.0, ge=0.0, le=1.0, description="Deterministic task score.")
+    matched_items: List[str] = Field(default_factory=list, description="Correctly matched items this step.")
+    missing_items: List[str] = Field(default_factory=list, description="Expected items still missing at grading time.")
+    notes: List[str] = Field(default_factory=list, description="Human-readable reward rationale.")
+
+
+class ClinicalTrialAction(Action):
+    """Agent action for extracting evidence and making screening decisions."""
+
+    action_type: Literal[
+        "extract_data",
+        "rank_patients",
+        "flag_deviation",
+        "submit_decision",
+        "delete_evidence",
+    ] = Field(..., description="Type of environment action.")
+    field_name: Optional[str] = Field(default=None, description="Clinical field being extracted.")
+    value: Optional[str] = Field(default=None, description="Normalized value for the extracted field.")
+    patient_id: Optional[str] = Field(default=None, description="Patient identifier for ranking or extraction.")
+    ranking: List[str] = Field(default_factory=list, description="Ordered patient IDs from best to worst fit.")
+    deviations: List[str] = Field(default_factory=list, description="Protocol deviations or exclusions identified.")
+    final_decision: Optional[str] = Field(
+        default=None,
+        description="Terminal decision such as eligible, ineligible, or ranking_submitted.",
+    )
+    rationale: Optional[str] = Field(default=None, description="Optional short reasoning trace.")
+
+
+class ClinicalTrialObservation(Observation):
+    """Observation returned after each environment interaction."""
+
+    task_id: str = Field(..., description="Current task identifier.")
+    difficulty: Literal["easy", "medium", "hard"] = Field(..., description="Task difficulty.")
+    title: str = Field(..., description="Scenario title.")
+    instructions: str = Field(..., description="Task instructions for the agent.")
+    context: Dict[str, Any] = Field(default_factory=dict, description="Structured and unstructured patient context.")
+    expected_fields: List[str] = Field(default_factory=list, description="High-value clinical fields to extract.")
+    extracted_fields: Dict[str, str] = Field(default_factory=dict, description="Validated data extracted so far.")
+    identified_deviations: List[str] = Field(default_factory=list, description="Validated protocol deviations found.")
+    attempts_remaining: int = Field(default=0, description="Steps left in the current episode.")
+    grader_name: str = Field(default="", description="Deterministic grader assigned to the task.")
+    reward_details: ClinicalTrialReward = Field(
+        default_factory=ClinicalTrialReward,
+        description="Structured reward breakdown for the step.",
+    )
+    terminal_reason: Optional[str] = Field(default=None, description="Why the episode ended, if done.")
+
+
+class ClinicalTrialState(State):
+    """Internal environment state exposed through the OpenEnv state endpoint."""
+
+    current_task_id: Optional[str] = Field(default=None, description="Current task identifier.")
+    difficulty: Optional[str] = Field(default=None, description="Difficulty for the current task.")
+    title: Optional[str] = Field(default=None, description="Current task title.")
+    extracted_fields: Dict[str, str] = Field(default_factory=dict, description="Accepted extracted fields.")
+    identified_deviations: List[str] = Field(default_factory=list, description="Accepted deviations.")
+    final_decision: Optional[str] = Field(default=None, description="Submitted terminal decision.")
+    grading_score: float = Field(default=0.0, ge=0.0, le=1.0, description="Latest grader output.")
+
+    def __call__(self) -> "ClinicalTrialState":
+        """Support env.state() as a compatibility alias for the OpenEnv state property."""
+        return self

openenv.yaml

@@ -0,0 +1,19 @@
+spec_version: 1
+version: 1.0.0
+name: clinical_trial_env
+author: Abhishek-CS221006
+entry_point: env:ClinicalTrialEnv
+tasks:
+  - id: easy
+    name: Metastatic NSCLC Screening
+    difficulty: easy
+  - id: medium
+    name: Patient Protocol Ranking
+    difficulty: medium
+  - id: hard
+    name: AML Exclusion Detection
+    difficulty: hard
+tags:
+  - healthcare
+  - clinical-trials
+  - medical-nlp

patient_data.json

@@ -0,0 +1,143 @@
+{
+  "tasks": {
+    "easy": {
+      "difficulty": "easy",
+      "title": "EGFR-Mutated Metastatic NSCLC Eligibility",
+      "instructions": "Review the structured oncology chart and determine whether the patient meets the five binary eligibility criteria for a first-line EGFR inhibitor trial. Use extract_data for clinical facts and submit_decision with eligible or ineligible.",
+      "context": {
+        "trial_protocol": {
+          "phase": "Phase III",
+          "disease": "Metastatic non-small cell lung cancer",
+          "binary_criteria": [
+            "Age >= 18 years",
+            "Pathologically confirmed metastatic NSCLC",
+            "Activating EGFR mutation present",
+            "ECOG performance status 0 or 1",
+            "No prior EGFR-targeted therapy"
+          ]
+        },
+        "patient": {
+          "patient_id": "NSCLC-204",
+          "age": 56,
+          "sex": "female",
+          "diagnosis": "Stage IV lung adenocarcinoma with pleural metastases",
+          "molecular_profile": "EGFR L858R positive; ALK negative; ROS1 negative",
+          "ecog": 1,
+          "prior_therapy": "Carboplatin/pemetrexed deferred; no prior EGFR TKI exposure",
+          "recent_mri": "No untreated brain metastases"
+        }
+      },
+      "ground_truth": {
+        "extracted_fields": {
+          "age": "56",
+          "diagnosis": "metastatic nsclc",
+          "egfr_mutation": "l858r positive",
+          "ecog": "1",
+          "prior_egfr_tki": "none"
+        },
+        "ranking": [],
+        "final_decision": "eligible"
+      },
+      "hidden_exclusions": [],
+      "max_steps": 6,
+      "grader_name": "deterministic_json_grader"
+    },
+    "medium": {
+      "difficulty": "medium",
+      "title": "Rank Patients for HER2-Positive Breast Cancer Trial",
+      "instructions": "Rank the three candidates from best to worst fit for a second-line HER2-positive metastatic breast cancer antibody-drug conjugate trial. Extract key biomarkers or performance status if useful, then submit the full ranking with rank_patients.",
+      "context": {
+        "trial_protocol": {
+          "phase": "Phase II",
+          "disease": "HER2-positive metastatic breast cancer",
+          "fit_score_drivers": [
+            "HER2 IHC 3+ or ISH amplified",
+            "Prior trastuzumab and taxane exposure required",
+            "No active CNS progression",
+            "ECOG 0-1 preferred",
+            "Adequate left ventricular ejection fraction"
+          ]
+        },
+        "patients": [
+          {
+            "patient_id": "BC-101",
+            "age": 48,
+            "her2_status": "IHC 3+",
+            "prior_therapy": "Taxane, trastuzumab, pertuzumab completed",
+            "ecog": 0,
+            "brain_mri": "Stable treated cerebellar lesion; no active CNS progression",
+            "lvef": "61%"
+          },
+          {
+            "patient_id": "BC-103",
+            "age": 62,
+            "her2_status": "ISH amplified",
+            "prior_therapy": "Taxane and trastuzumab completed; pertuzumab intolerant",
+            "ecog": 1,
+            "brain_mri": "No CNS disease",
+            "lvef": "55%"
+          },
+          {
+            "patient_id": "BC-102",
+            "age": 54,
+            "her2_status": "IHC 2+ / ISH non-amplified",
+            "prior_therapy": "Taxane only; trastuzumab naive",
+            "ecog": 2,
+            "brain_mri": "New enhancing frontal lesion on last scan",
+            "lvef": "49%"
+          }
+        ]
+      },
+      "ground_truth": {
+        "extracted_fields": {
+          "BC-101_her2_status": "ihc 3+",
+          "BC-101_ecog": "0",
+          "BC-103_her2_status": "ish amplified",
+          "BC-103_ecog": "1",
+          "BC-102_cns_status": "active cns progression",
+          "BC-102_trastuzumab_exposure": "none"
+        },
+        "ranking": ["BC-101", "BC-103", "BC-102"],
+        "final_decision": "ranking_submitted"
+      },
+      "hidden_exclusions": [],
+      "max_steps": 6,
+      "grader_name": "deterministic_json_grader"
+    },
+    "hard": {
+      "difficulty": "hard",
+      "title": "Protocol Deviations in AML Screening Note",
+      "instructions": "Read the unstructured hematology note, identify protocol deviations or exclusions, and determine if the patient is screen-failed. Use extract_data for critical facts, flag_deviation for protocol violations, then submit_decision with eligible or ineligible.",
+      "context": {
+        "trial_protocol": {
+          "phase": "Phase Ib",
+          "disease": "Relapsed/refractory FLT3-mutated AML",
+          "key_exclusions": [
+            "Neutropenic fever / active uncontrolled infection",
+            "QTc > 480 ms",
+            "Strong CYP3A4 inhibitor within 7 days"
+          ]
+        },
+        "screening_note": "Mr. J.R. is a 63-year-old man with relapsed AML harboring FLT3-ITD. He presented overnight with neutropenic fever and positive blood cultures. ECG today showed QTc 486 ms. Posaconazole prophylaxis was restarted five days ago because of prior pulmonary aspergillosis. Bone marrow confirms persistent 22% blasts. ECOG 1. Creatinine clearance 68 mL/min."
+      },
+      "ground_truth": {
+        "extracted_fields": {
+          "age": "63",
+          "biomarker": "flt3-itd",
+          "ecg_qtc_ms": "486",
+          "infection_status": "neutropenic fever",
+          "cyp3a4_inhibitor": "posaconazole within 7 days"
+        },
+        "ranking": [],
+        "final_decision": "ineligible"
+      },
+      "hidden_exclusions": [
+        "neutropenic fever",
+        "qtc greater than 480 ms",
+        "recent strong cyp3a4 inhibitor"
+      ],
+      "max_steps": 6,
+      "grader_name": "deterministic_json_grader"
+    }
+  }
+}

pyproject.toml

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-clinical_trial_env"
+version = "0.1.0"
+description = "Clinical Trial Env environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    # install from github
+    # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
+    "openenv-core[core]>=0.2.2",
+    # Environment-specific dependencies
+    # Add all dependencies needed for your environment here
+    # Examples:
+    # "numpy>=1.19.0",
+    # "torch>=2.0.0",
+    # "gymnasium>=0.29.0",
+    # "openspiel>=1.0.0",
+    # "smolagents>=1.22.0,<2",
+    "openai>=1.30.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m clinical_trial_env.server.app
+server = "clinical_trial_env.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["clinical_trial_env", "clinical_trial_env.server"]
+package-dir = { "clinical_trial_env" = ".", "clinical_trial_env.server" = "server" }

server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Clinical Trial Env environment server components."""
+
+from .clinical_trial_env_environment import ClinicalTrialEnvironment
+
+__all__ = ["ClinicalTrialEnvironment"]

server/app.py

@@ -0,0 +1,79 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FastAPI application for the Clinical Trial Env Environment.
+
+This module creates an HTTP server that exposes the ClinicalTrialEnvironment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment
+    - POST /step: Execute an action
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m server.app
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
+    ) from e
+
+try:
+    from ..models import ClinicalTrialAction, ClinicalTrialObservation
+    from .clinical_trial_env_environment import ClinicalTrialEnvironment
+except ImportError:
+    from models import ClinicalTrialAction, ClinicalTrialObservation
+    from server.clinical_trial_env_environment import ClinicalTrialEnvironment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+    ClinicalTrialEnvironment,
+    ClinicalTrialAction,
+    ClinicalTrialObservation,
+    env_name="clinical_trial_env",
+    max_concurrent_envs=4,
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    This function enables running the server without Docker:
+        uv run --project . server
+        uv run --project . server --port 8001
+        python -m clinical_trial_env.server.app
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+
+    For production deployments, consider using uvicorn directly with
+    multiple workers:
+        uvicorn clinical_trial_env.server.app:app --workers 4
+    """
+    import uvicorn
+
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()

server/clinical_trial_env_environment.py

@@ -0,0 +1,10 @@
+"""Compatibility wrapper for the root clinical trial environment."""
+
+from __future__ import annotations
+
+try:
+    from ..env import ClinicalTrialEnvironment
+except ImportError:
+    from env import ClinicalTrialEnvironment
+
+__all__ = ["ClinicalTrialEnvironment"]

server/requirements.txt

@@ -0,0 +1,6 @@
+openenv[core]>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+
+
+