Version 3

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.gitignore

@@ -23,9 +23,13 @@ yarn-debug.log*
 yarn-error.log*
 pnpm-debug.log*
 
+# --- Project-specific ---
+PROMPT.md
+
 # --- Build / temp ---
 *.log
 *.tmp
 *.swp
 .DS_Store
 
+arXiv-2212.05190v3/

Dockerfile

@@ -1,7 +1,7 @@
 FROM node:20-alpine AS frontend-builder
 WORKDIR /app/frontend
 COPY frontend/package*.json ./
-RUN npm ci
+RUN npm ci --production=false
 COPY frontend/ ./
 RUN npm run build
 
@@ -11,6 +11,8 @@ RUN apt-get update && \
     apt-get install -y --no-install-recommends build-essential curl && \
     rm -rf /var/lib/apt/lists/*
 
+# HF Spaces runs as uid 1000
+RUN useradd -m -u 1000 user
 WORKDIR /app
 
 COPY backend/requirements.txt /app/backend/requirements.txt
@@ -22,18 +24,26 @@ COPY scripts /app/scripts
 COPY openenv.yaml /app/openenv.yaml
 COPY .env.example /app/.env.example
 COPY inference.py /app/inference.py
+COPY train_rl.py /app/train_rl.py
+COPY train_bandit.py /app/train_bandit.py
 
 COPY --from=frontend-builder /app/frontend/dist /app/frontend/dist
 
 RUN python3 /app/scripts/preprocess_data.py
 
+# Ensure the user owns the app directory and has a writable home (HF Spaces)
+RUN chown -R user:user /app && \
+    mkdir -p /home/user/.cache && chown -R user:user /home/user
+
 ENV PORT=7860
-ENV PYTHONPATH="/app/backend/src:${PYTHONPATH}"
+ENV PYTHONPATH="/app/backend/src:/app"
 ENV PYTHONUNBUFFERED=1
 
+USER user
+
 EXPOSE 7860
 
-HEALTHCHECK --interval=30s --timeout=3s --start-period=15s --retries=3 \
+HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \
     CMD curl -f http://localhost:7860/health || exit 1
 
-CMD ["sh", "-c", "uvicorn backend.main:app --host 0.0.0.0 --port ${PORT:-7860}"]
+CMD ["sh", "-c", "uvicorn backend.main:app --host 0.0.0.0 --port ${PORT:-7860} --workers 1"]

PROMPT.md

@@ -1,571 +0,0 @@
-You are an expert Python backend, ML, and infrastructure engineer.
-Your task is to implement a complete, production-ready OpenEnv environment called **PolypharmacyEnv** for training and evaluating agentic RL policies that act as an "elderly polypharmacy safety agent" (clinical pharmacist assistant).
-
-The deliverable MUST satisfy all of the following:
-- Fully compliant with the OpenEnv spec (typed models, `step()` / `reset()` / `state()`, `openenv.yaml`, HTTP server, Dockerfile).
-- Simulates a realistic healthcare workflow around elderly polypharmacy and dangerous drug combinations.
-- Defines at least **3 tasks** (easy → medium → hard) with deterministic agent graders producing scores in (0.0, 1.0).
-- Provides shaped rewards over the trajectory (not just sparse terminal rewards).
-- Includes a baseline LLM-based inference script `inference.py` in the repo root, following the evaluation requirements:
-  - Uses the OpenAI Python client.
-  - Reads `OPENAI_API_KEY`, `API_BASE_URL`, `MODEL_NAME`, and `HF_TOKEN` from the environment.
-  - Emits structured stdout logs in the exact `[START]`, `[STEP]`, `[END]` format from the OpenEnv sample inference script.
-- Is containerized and deployable as a **Hugging Face Space** tagged with `openenv` that responds to OpenEnv-style `reset` / `step` / `state` HTTP calls.
-
-Implement everything described below.
-
-=================================================
-1. Repository and folder structure
-=================================================
-
-Create a Python package repository with this structure (names are important unless clearly labeled as examples):
-
-- `openenv-polypharmacy/`
-  - `openenv.yaml`
-  - `README.md`
-  - `requirements.txt`
-  - `Dockerfile`
-  - `inference.py`                 # baseline LLM agent per spec
-  - `pyproject.toml` or `setup.cfg` (optional but recommended)
-  - `src/`
-    - `polypharmacy_env/`
-      - `__init__.py`
-      - `config.py`
-      - `models.py`                # Action, Observation, State, helper models
-      - `env_core.py`              # PolypharmacyEnv implementation
-      - `tasks.py`                 # task setup utilities
-      - `graders.py`               # deterministic graders for each task
-      - `rewards.py`               # reward shaping logic
-      - `data_loader.py`           # load/preprocess patient and lookup data
-      - `ddi_simulator.py`         # local DDI / guideline simulator
-      - `api/`
-        - `__init__.py`
-        - `schemas.py`            # HTTP request/response schemas
-        - `server.py`             # FastAPI app exposing OpenEnv endpoints
-      - `baselines/`
-        - `__init__.py`
-        - `heuristic_agent.py`    # simple rule-based baseline agent
-        - `random_agent.py`       # trivial random baseline (optional)
-      - `tests/`
-        - `__init__.py`
-        - `test_env_core.py`
-        - `test_api.py`
-  - `data/`
-    - `raw/`                      # placeholder for real/synthetic source data
-    - `processed/`
-    - `lookups/`
-      - `ddi_rules.csv`
-      - `beers_criteria.csv`
-      - `drug_metadata.csv`
-  - `scripts/`
-    - `preprocess_data.py`
-    - `run_validation.sh`         # optional; runs OpenEnv validator, tests, etc.
-
-Use Python 3.10+ with full type hints, and keep the code black/isort-compatible.
-
-=================================================
-2. Domain, data, and clinical abstraction
-=================================================
-
-2.1. Core scenario
-
-Model an elderly patient (age ≥ 65) with:
-- Demographics: age, sex.
-- Comorbidities: e.g., hypertension, diabetes, heart failure, CKD, dementia.
-- Basic labs: kidney function (eGFR category), liver function category.
-- A current medication list (polypharmacy, e.g., 3–15 drugs depending on task).
-
-Each **episode** is one medication-review session where the agent:
-- Observes patient info and current meds.
-- Optionally **queries** a DDI/guideline tool for specific drug pairs.
-- Proposes **interventions**:
-  - `stop`: discontinue a drug.
-  - `dose_reduce`: lower dose of a drug.
-  - `substitute`: swap to a safer alternative.
-  - `add_monitoring`: keep the drug but flag extra monitoring.
-- Calls `finish_review` when it decides the regimen is acceptable or budgets are exhausted.
-
-No external PHI, EHRs, or online APIs: all data is **synthetic** or de-identified and local to the container (CSV files).
-
-2.2. Data files and CSV schemas
-
-Implement local CSVs under `data/lookups/`:
-
-**`drug_metadata.csv`**
-- `drug_id` (string; unique key)
-- `generic_name` (string)
-- `atc_class` (string)
-- `is_high_risk_elderly` (0/1)
-- `default_dose_mg` (float)
-- `min_dose_mg` (float)
-- `max_dose_mg` (float)
-
-**`beers_criteria.csv`**
-- `drug_id` (string)
-- `criterion_type` (enum string: `avoid`, `caution`, `dose_adjust`, `avoid_in_condition`)
-- `condition` (nullable string; e.g., `CKD`, `dementia`)
-- `rationale` (brief text)
-
-**`ddi_rules.csv`**
-- `drug_id_1` (string; normalized so `drug_id_1 < drug_id_2` lexicographically)
-- `drug_id_2` (string)
-- `severity` (enum string: `mild`, `moderate`, `severe`)
-- `mechanism` (short text)
-- `recommendation` (enum string: `avoid_combination`, `monitor_closely`, `dose_adjust`, `no_action`)
-- `base_risk_score` (float in [0.0, 1.0])
-
-Implement a synthetic patient-episode dataset under `data/processed/`:
-
-**`patients_polypharmacy.csv`**
-- `episode_id` (string)
-- `age` (int)
-- `sex` (enum: `M`, `F`, `O`)
-- `conditions` (semicolon-separated; e.g., `HTN;DM;CKD`)
-- `eGFR_category` (enum: `normal`, `mild`, `moderate`, `severe`)
-- `liver_function_category` (enum: `normal`, `impaired`)
-- `medication_ids` (semicolon-separated list of `drug_id`)
-- `baseline_risk_score` (float in [0.0, 1.0])
-
-2.3. Preprocessing script
-
-In `scripts/preprocess_data.py`:
-- If real data is not provided, procedurally generate synthetic but plausible data using:
-  - Random combinations of conditions and drugs constrained by simple rules (e.g., CKD + renally-cleared drugs).
-  - Controlled distribution of high-risk DDIs and Beers violations.
-- Explicitly tag episodes as easy/medium/hard (e.g., via number of drugs, number/severity of DDIs, and number of Beers issues).
-- Save `patients_polypharmacy.csv` ready for the environment to consume.
-
-=================================================
-3. OpenEnv models and environment implementation
-=================================================
-
-3.1. Models
-
-In `models.py`, define dataclasses or Pydantic models that extend the appropriate OpenEnv base types (`Action`, `Observation`, `State`) and are JSON-compatible.
-
-Auxiliary models:
-
-**`MedicationEntry`**
-- `drug_id: str`
-- `generic_name: str`
-- `atc_class: str`
-- `dose_mg: float`
-- `frequency: str`          # e.g., `qd`, `bid`
-- `route: str`              # e.g., `po`
-- `is_high_risk_elderly: bool`
-- `beers_flags: list[str]`  # e.g., `["avoid", "dose_adjust_CKD"]`
-
-**`InteractionQueryRecord`**
-- `drug_id_1: str`
-- `drug_id_2: str`
-- `severity: str | None`
-- `recommendation: str | None`
-- `risk_score: float | None`
-- `step_index: int`
-
-**`InterventionRecord`**
-- `target_drug_id: str`
-- `action_type: Literal["stop", "dose_reduce", "substitute", "add_monitoring"]`
-- `proposed_new_drug_id: str | None`
-- `rationale: str`
-- `step_index: int`
-
-Core wire models:
-
-**`PolypharmacyObservation`** (extends OpenEnv `Observation`)
-- `episode_id: str`
-- `task_id: Literal["easy_screening", "budgeted_screening", "complex_tradeoff"]`
-- `age: int`
-- `sex: str`
-- `conditions: list[str]`
-- `eGFR_category: str`
-- `liver_function_category: str`
-- `current_medications: list[MedicationEntry]`
-- `interaction_queries: list[InteractionQueryRecord]`
-- `interventions: list[InterventionRecord]`
-- `step_index: int`
-- `remaining_query_budget: int`
-- `remaining_intervention_budget: int`
-- `shaped_reward: float`  # reward from last step
-- `done: bool`
-
-**`PolypharmacyAction`** (extends OpenEnv `Action`)
-- `action_type: Literal["query_ddi", "propose_intervention", "finish_review"]`
-- `drug_id_1: str | None`        # for DDI queries or some interventions
-- `drug_id_2: str | None`        # for DDI queries
-- `target_drug_id: str | None`   # for interventions
-- `intervention_type: Literal["stop", "dose_reduce", "substitute", "add_monitoring", "none"] | None`
-- `proposed_new_drug_id: str | None`
-- `rationale: str | None`
-
-**`PolypharmacyState`** (extends OpenEnv `State`)
-- `episode_id: str`
-- `task_id: str`
-- `step_count: int`
-- `max_steps: int`
-- `num_query_actions: int`
-- `num_interventions: int`
-
-3.2. Environment core
-
-In `env_core.py`, implement `PolypharmacyEnv` extending the appropriate OpenEnv environment base class. It must implement:
-
-**`reset(task_id: str | None = None) -> PolypharmacyObservation`**
-- If `task_id` is `None`, default to medium (`budgeted_screening`).
-- Sample an episode from `patients_polypharmacy.csv` filtered by difficulty.
-- Initialize:
-  - `episode_id`
-  - `step_count = 0`
-  - task-specific budgets (query, interventions, max_steps)
-  - baseline regime and risk
-  - empty `interaction_queries` and `interventions`
-- Return the initial `PolypharmacyObservation` with:
-  - `step_index = 0`
-  - `shaped_reward = 0.0`
-  - `done = False`
-
-**`step(action: PolypharmacyAction) -> dict`**
-- Validate the action; if invalid:
-  - Apply a negative reward.
-  - Do not modify regimen, but log error in `info`.
-- If `action_type == "query_ddi"`:
-  - If query budget exhausted, apply penalty and do not query.
-  - Else:
-    - Use `ddi_simulator.lookup_ddi(drug_id_1, drug_id_2)` to get severity, recommendation, base_risk_score.
-    - Append an `InteractionQueryRecord`.
-    - Apply a small negative reward for query cost.
-- If `action_type == "propose_intervention"`:
-  - If intervention budget exhausted, apply penalty and ignore change.
-  - Else:
-    - Update `current_medications` according to `intervention_type`:
-      - `stop`: remove medication.
-      - `dose_reduce`: adjust dose downward within [min_dose_mg, default_dose_mg].
-      - `substitute`: replace with a safer alternative from same `atc_class`.
-      - `add_monitoring`: keep drug but tag in internal state.
-    - Append an `InterventionRecord`.
-    - Recompute current regimen risk using the risk model (see 3.3).
-    - Compute shaped reward = (previous_risk - new_risk) - small intervention cost.
-- If `action_type == "finish_review"`:
-  - Mark `done = True`.
-  - Call the task’s grader to get episode-level score in [0.0, 1.0].
-  - Add this as a terminal bonus to the current step reward.
-
-- In all cases:
-  - Increment `step_count`.
-  - Check `max_steps`; if exceeded, auto-terminate:
-    - `done = True`
-    - apply time-out penalty
-    - call grader with current trajectory for a final score if appropriate.
-  - Construct next `PolypharmacyObservation` with updated fields.
-  - Return a dict:
-    - `observation`: `PolypharmacyObservation`
-    - `reward`: float shaped reward for this step
-    - `done`: bool
-    - `info`: dict with fields like `current_risk`, `baseline_risk`, `grader_score_if_terminal`, and debug flags.
-
-**`state` property**
-- Returns `PolypharmacyState` reflecting the current internal state.
-
-3.3. DDI simulator and risk model
-
-In `ddi_simulator.py`:
-- Load `ddi_rules.csv` once via `data_loader`.
-- Implement `lookup_ddi(drug_id_1, drug_id_2) -> tuple[severity, recommendation, base_risk_score]`:
-  - Normalize the pair ordering.
-  - Look up row; if missing, return:
-    - severity = `"none"`
-    - recommendation = `"no_action"`
-    - base_risk_score = 0.0
-
-In `rewards.py` (or a dedicated module), implement:
-- `compute_regimen_risk(current_drug_ids, patient_context, ddi_rules, beers_rules, drug_metadata) -> float`
-  - Aggregate contributions from:
-    - Beers violations (weighted by `criterion_type` and relevant conditions).
-    - DDI base risk scores for all present drug pairs.
-    - High-risk elderly drugs.
-  - Normalize and clip to [0.0, 1.0].
-
-Use this function to compute:
-- `baseline_risk` at episode start.
-- Risk after each intervention step.
-
-Also implement:
-- `compute_shaped_reward(previous_risk, new_risk, action, context, partial_metrics) -> float`
-  - Positive component: `previous_risk - new_risk`.
-  - Negative components: per-query cost, per-intervention cost, invalid-action penalty, time-out penalty.
-
-=================================================
-4. Tasks and graders (3 difficulty levels)
-=================================================
-
-Define three task IDs and semantics in `tasks.py` and `graders.py`:
-
-Task IDs:
-- `easy_screening`
-- `budgeted_screening`
-- `complex_tradeoff`
-
-4.1. `easy_screening` (easy)
-
-- Small regimen: 3–5 drugs.
-- Exactly one **severe** DDI pair and possibly one simple Beers violation.
-- Budgets:
-  - query_budget ≈ 4
-  - intervention_budget ≈ 2
-  - max_steps ≈ 10
-
-Grader:
-- Input: full trajectory, baseline risk, final risk, list of interventions.
-- Compute:
-  - `risk_reduction = max(0.0, baseline_risk - final_risk) / max(baseline_risk, ε)` (normalized).
-  - `targeted_intervention_flag = 1.0` if at least one intervention affects one of the drugs in the known severe DDI pair, else 0.0.
-- Score:
-  - `score = 0.5 * risk_reduction + 0.5 * targeted_intervention_flag`
-  - Clip to [0.0, 1.0].
-
-4.2. `budgeted_screening` (medium)
-
-- Medium regimen: 6–10 drugs.
-- Multiple DDIs (mild/moderate/severe) and multiple Beers issues.
-- Budgets:
-  - query_budget ≈ 8
-  - intervention_budget ≈ 3
-  - max_steps ≈ 20
-
-Grader:
-- Compute:
-  - `risk_reduction_score` as normalized risk drop.
-  - `intervention_precision_score` = fraction of interventions that actually reduce risk or fix guideline violations.
-  - `query_efficiency_score` = (number of severe/moderate DDIs discovered) / (number of queries used), normalized.
-- Weighted score, for example:
-  - `score = 0.5 * risk_reduction_score + 0.3 * intervention_precision_score + 0.2 * query_efficiency_score`
-  - Clip to [0.0, 1.0].
-
-4.3. `complex_tradeoff` (hard)
-
-- Larger regimen: 10–15 drugs.
-- Some drugs are **clinically critical** (e.g., anticoagulants, insulin analogues) and encoded as such in `drug_metadata` or a small internal map.
-- Episodes contain:
-  - multiple DDIs and Beers issues, including ones involving critical drugs.
-  - safer substitutes for some risky drugs.
-
-Budgets:
-- query_budget ≈ 12
-- intervention_budget ≈ 5
-- max_steps ≈ 30
-
-Grader adds a **regimen disruption penalty** component:
-- Metrics:
-  - `risk_reduction_score` (as above).
-  - `critical_drug_penalty` = penalty if a critical drug is stopped without substitution to another suitable agent.
-  - `total_drug_changes` = number of drugs stopped or substituted.
-  - `regimen_disruption_penalty` derived from `total_drug_changes` and `critical_drug_penalty`.
-
-Example scoring:
-- `base = risk_reduction_score`
-- `penalty = α * regimen_disruption_penalty`
-- `score = clamp(base - penalty, 0.0, 1.0)`
-
-4.4. Reward shaping
-
-In `rewards.py`, define a consistent shaping scheme:
-- On each query:
-  - Small negative reward (e.g., −0.01) plus any small bonus if it discovers a severe DDI, if desired.
-- On each intervention:
-  - Reward ≈ (previous_risk - new_risk) − small intervention cost.
-- On invalid actions:
-  - Larger negative reward (e.g., −0.1) and no state change.
-- On `finish_review`:
-  - Add the task-level `score` ∈ [0.0, 1.0] from the corresponding grader to that step’s shaped reward.
-
-Ensure the sum of step rewards per episode remains in a reasonable numeric range (e.g., roughly -5 to +5) while still allowing meaningful differentiation by graders.
-
-=================================================
-5. HTTP API server and openenv.yaml
-=================================================
-
-5.1. HTTP server (FastAPI)
-
-In `api/server.py`:
-- Implement a FastAPI app that maintains a `PolypharmacyEnv` instance (or a multiplexing scheme if needed).
-- Endpoints:
-  - `POST /reset`:
-    - Request body: may include `task_id` (string).
-    - Response: serialized `PolypharmacyObservation`.
-  - `POST /step`:
-    - Request body: serialized `PolypharmacyAction`.
-    - Response: dict with:
-      - `observation`: `PolypharmacyObservation`
-      - `reward`: float
-      - `done`: bool
-      - `info`: dict
-  - `GET /state`:
-    - Response: `PolypharmacyState`.
-
-Provide a module-level `app = FastAPI(...)` object for use with uvicorn and Hugging Face Spaces. Ensure the JSON schema is consistent with OpenEnv clients (simple, flat JSON for observation/action/state).
-
-5.2. `openenv.yaml`
-
-At repo root, define `openenv.yaml` consistent with the latest OpenEnv spec. At minimum, include:
-- `name`: `polypharmacy_env`
-- `version`: e.g., `0.1.0`
-- `description`: human-readable description.
-- `author`: your details.
-- `tags`: e.g., `["healthcare", "polypharmacy", "openenv"]`
-- `tasks`:
-  - One entry per task:
-    - `id`: `"easy_screening"` / `"budgeted_screening"` / `"complex_tradeoff"`
-    - `description`: one-line description
-    - `difficulty`: `"easy"`, `"medium"`, `"hard"`
-
-Ensure `openenv validate` (or equivalent validator) passes once implemented.
-
-=================================================
-6. Baseline heuristic (non-LLM) agent
-=================================================
-
-In `baselines/heuristic_agent.py`, implement a simple, deterministic baseline agent that:
-
-For each episode:
-- Iterates through all unordered medication pairs within query budget:
-  - Calls `query_ddi` via the environment for each pair until the query budget is exhausted or all pairs are examined.
-  - Records severe and moderate interactions.
-- After querying:
-  - For each severe DDI pair:
-    - Try `substitute` one of the drugs using `drug_metadata`:
-      - Prefer substitute within same `atc_class` that:
-        - is not marked high-risk elderly.
-        - does not participate in known severe DDIs with the rest of the regimen.
-    - If no substitute exists, propose `stop` for the higher-risk drug.
-  - Respect intervention budget limits.
-- Finally, call `finish_review`.
-
-This baseline should be callable as a simple Python function that interacts with `PolypharmacyEnv` directly (without HTTP).
-
-=================================================
-7. Baseline LLM inference script (inference.py)
-=================================================
-
-At repo root, create `inference.py` that:
-
-7.1. Uses the OpenAI Python client
-
-- Import and configure the official OpenAI Python client.
-- Read environment variables:
-  - `OPENAI_API_KEY` (required).
-  - `API_BASE_URL` (base URL for LLM; default to OpenAI standard if not set).
-  - `MODEL_NAME` (e.g., `gpt-4.1` or similar).
-  - `HF_TOKEN` (if needed for HF auth; do not hardcode).
-- Read `POLYPHARMACY_ENV_URL` (or similar) for the environment’s HTTP base URL.
-
-7.2. Implements the required logging format
-
-- For each **run** across all tasks:
-  - Emit a `[START]` line with a JSON payload exactly matching the evaluation specification:
-    - Fields such as `run_id`, `task_id`, `model`, etc., in the same order and naming as the sample OpenEnv inference script.
-- For each **step** in an episode:
-  - Emit a `[STEP]` line with JSON fields including:
-    - `run_id`
-    - `task_id`
-    - `episode_id`
-    - `step_index`
-    - `observation_summary` (brief, machine-readable summary)
-    - `action_payload` (the action sent to the env)
-    - `reward`
-    - `done`
-- After finishing an episode for a task:
-  - Emit an `[END]` line summarizing:
-    - `run_id`
-    - `task_id`
-    - per-episode statistics (e.g., total reward, grader score from last step’s `info`).
-- The stdout format MUST follow the sample exactly:
-  - Same tags: `[START]`, `[STEP]`, `[END]`.
-  - Same JSON field names and ordering as the provided reference.
-  - No extra prints except these structured logs (and necessary error messages to stderr).
-
-7.3. LLM agent loop
-
-- For each task (`easy_screening`, `budgeted_screening`, `complex_tradeoff`):
-  - Run a fixed small number of episodes (e.g., 5–10 per task) for baseline scoring.
-  - For each episode:
-    - Call `/reset` with the task id.
-    - At each step:
-      - Summarize the observation into a concise prompt for the LLM:
-        - Include age, sex, conditions, high-risk flags, budgets, and a compressed view of meds and previous actions.
-      - Ask the model to output a **strict JSON** representing `PolypharmacyAction` fields.
-      - Parse and validate the JSON; if invalid, fall back to a safe default (e.g., `finish_review` or a no-op) and penalize in evaluation.
-      - Send this action to `/step` and log `[STEP]`.
-    - End when `done=True` or max_steps is reached.
-- At the end, print aggregate scores per task and overall.
-
-Make sure runtime < 20 minutes and that the script can run within 2 vCPUs and 8 GB RAM.
-
-=================================================
-8. Dockerfile and Hugging Face Space
-=================================================
-
-8.1. Dockerfile
-
-Create a `Dockerfile` that:
-- Starts from a slim Python image (e.g., `python:3.11-slim`).
-- Installs system dependencies as needed (e.g., `build-essential`, `curl`).
-- Copies the project into the container.
-- Installs Python dependencies from `requirements.txt`.
-- Sets appropriate environment variables for the app (e.g., `PORT=7860`).
-- Exposes port 7860.
-- Uses a `CMD` or `ENTRYPOINT` that runs the FastAPI server, for example:
-  - `uvicorn polypharmacy_env.api.server:app --host 0.0.0.0 --port 7860`
-
-8.2. Hugging Face Space
-
-Ensure the repository is ready to be used as a Hugging Face Space:
-- Space type: `docker`.
-- Tag: `openenv`.
-- On container start, the server must listen on the correct port and respond to:
-  - `POST /reset`
-  - `POST /step`
-  - `GET /state`
-- The environment must start cleanly with `docker build` + `docker run` locally.
-
-=================================================
-9. README and documentation
-=================================================
-
-In `README.md`, include:
-
-- **Environment description & motivation**:
-  - What PolypharmacyEnv simulates.
-  - Why elderly polypharmacy safety matters.
-- **Action and observation spaces**:
-  - Describe `PolypharmacyAction`, `PolypharmacyObservation`, and `PolypharmacyState` fields and semantics.
-- **Task descriptions**:
-  - `easy_screening`, `budgeted_screening`, `complex_tradeoff`, their difficulty and goals.
-- **Reward structure**:
-  - Summarize shaping and terminal rewards.
-- **Setup & usage**:
-  - How to install dependencies.
-  - How to run the API server locally (uvicorn command).
-  - How to run the heuristic baseline.
-  - How to run `inference.py` with environment variables.
-- **Baseline scores**:
-  - Document reproducible baseline scores for each task (heuristic agent, and LLM baseline if available).
-
-=================================================
-10. Validation and quality gates
-=================================================
-
-- Ensure:
-  - `openenv.yaml` and the HTTP server pass the OpenEnv validation script.
-  - `docker build` and `docker run` work without errors.
-  - `inference.py` completes under 20 minutes, within 2 vCPUs / 8 GB RAM.
-  - All graders:
-    - Are deterministic.
-    - Return scores strictly in [0.0, 1.0].
-  - No grader returns a constant score irrespective of behavior.
-
-Aim for clean, well-structured, well-documented code with clear separation of concerns between:
-- Data loading,
-- Environment state & dynamics,
-- Reward/grade logic,
-- HTTP serving,
-- Baseline agents and inference.

README.MD

@@ -1,256 +1,528 @@
 ---
-title: Polypharmacy
-emoji: 📉
-colorFrom: yellow
-colorTo: blue
+title: PolypharmacyEnv
+emoji: 💊
+colorFrom: blue
+colorTo: purple
 sdk: docker
 pinned: false
+tags:
+  - openenv
+  - healthcare
+  - reinforcement-learning
 ---
 
-# PolypharmacyEnv
+# PolypharmacyEnv — Elderly Medication Safety via Reinforcement Learning
 
-Monorepo for an OpenEnv-compatible medication safety environment with:
+An [OpenEnv](https://github.com/meta-pytorch/OpenEnv)-compliant environment that simulates **elderly polypharmacy medication review**. An RL agent acts as a clinical pharmacist assistant: it queries drug-drug interactions (DDIs), identifies Beers-criteria violations, and proposes safe interventions — all under resource-constrained budgets.
 
-- a FastAPI backend (`backend/`)
-- a React frontend (`frontend/`)
-- data assets (`data/`)
-- utility scripts (`scripts/`)
+Built for the **PyTorch OpenEnv Hackathon** to demonstrate how clinical decision support for polypharmacy can be framed as a sequential RL problem and served as a reusable environment through the OpenEnv hub.
+
+---
+
+## Why This Matters
+
+Polypharmacy — the simultaneous use of five or more medications — affects the majority of adults over 65. Elderly patients often see multiple specialists who may not be aware of each other's prescriptions, leading to dangerous drug combinations. Studies report that **adverse drug events from polypharmacy contribute to 100,000+ hospitalizations annually** in the US alone.
+
+Current solutions use static risk scoring. PolypharmacyEnv goes further by framing medication review as a **sequential decision problem**, where an RL agent must strategically allocate limited query and intervention budgets to maximize patient safety — exactly the kind of resource-constrained optimization that reinforcement learning excels at.
+
+**Reference**: Larouche, A., Durand, A., Khoury, R. & Sirois, C. (2023). [Neural Bandits for Data Mining: Searching for Dangerous Polypharmacy](https://link.springer.com/chapter/10.1007/978-3-031-36938-4_5). *Advances in Artificial Intelligence*, Springer.
+
+---
+
+## How OpenEnv & RL Power This
+
+### The RL Formulation
+
+PolypharmacyEnv frames medication review as a **Markov Decision Process (MDP)**:
+
+- **State**: Patient profile (age, conditions, organ function) + current medication list + interaction history
+- **Action space**: `query_ddi(drug_i, drug_j)` | `propose_intervention(target, type)` | `finish_review`
+- **Reward**: Shaped, dense signal at every step (not sparse end-of-episode). Queries cost budget (-0.015), discovering severe DDIs earns bonus (+0.05), successful interventions earn proportional risk reduction minus cost, invalid actions are penalized (-0.15), and `finish_review` triggers a grader that returns a terminal score in [0.0, 1.0].
+- **Constraint**: Finite query and intervention budgets, creating a resource-allocation optimization problem.
+
+This MDP is what makes the problem fundamentally different from static risk scoring: the agent must **decide what information to acquire** (which drug pairs to query) and **which interventions to prioritize**, all under budget constraints — a sequential decision problem that RL is designed to solve.
+
+### OpenEnv Interface
+
+PolypharmacyEnv implements the full **OpenEnv standard**:
+
+- **`reset()`** — Generates a new patient scenario (age, conditions, medication list)
+- **`step(action)`** — Processes an agent action, updates regimen state, returns shaped reward
+- **`state()`** — Returns the current episode snapshot
+
+All models use typed Pydantic classes extending OpenEnv base types (`PolypharmacyAction`, `PolypharmacyObservation`, `PolypharmacyState`).
+
+### What the Environment Enables
+
+The shaped reward function provides continuous signal over the full trajectory, making this environment compatible with standard RL training approaches:
+
+- **Policy gradient methods** (REINFORCE, PPO, GRPO): The per-step reward signal allows policy networks to learn query prioritization and intervention strategies.
+- **OpenEnv training pipeline**: Through OpenEnv's `step()`/`reset()` HTTP interface, external RL training loops can connect to this environment and train policies without modification.
+- **Neural Bandits (OptimNeuralTS)**: The budget-constrained query selection implements the OptimNeuralTS approach from the reference paper — Neural Thompson Sampling combined with Differential Evolution for efficient search.
+
+### Included Agents
+
+The repository ships with multiple agent implementations spanning rule-based, RL-trained, bandit-based, and LLM-based approaches:
+
+- **OptimNeuralTS bandit** (`train_bandit.py`, `neural_bandits.py`): Implements the paper's core algorithm — Neural Thompson Sampling with Differential Evolution to efficiently search for dangerous drug combinations. Builds an ensemble of models across training steps for high-precision predictions.
+- **REINFORCE-trained policy** (`train_rl.py`): A neural network policy trained via REINFORCE with learned baseline against the environment's shaped reward. Demonstrates that the MDP formulation and reward shaping enable genuine policy improvement through RL training.
+- **Heuristic agent** (`baselines/heuristic_agent.py`): Deterministic rule-based strategy that queries high-risk drug pairs first, then intervenes on severe DDIs. Serves as a strong domain-knowledge baseline.
+- **LLM agent** (`inference.py`): Uses an LLM (Qwen2.5-72B via OpenAI-compatible API) for zero-shot action generation. Demonstrates baseline LLM performance without RL fine-tuning.
+- **AI suggestion endpoint** (`/agent/suggest`): LLM-powered action suggestions with rule-based guardrails for the interactive UI.
 
 ---
 
 ## Repository Structure
 
-```text
-backend/
-    main.py                      # ASGI entrypoint (uvicorn target)
-    requirements.txt             # Backend dependencies
-    Dockerfile                   # Backend container
-    src/polypharmacy_env/        # Python package source
-      api/
-        app.py                   # FastAPI/OpenEnv app assembly
-        server.py                # Compatibility import wrapper
-        routes/agent.py          # /agent/suggest route
-      services/
-        groq_agent.py            # Groq-based action suggestion logic
-      env_core.py                # OpenEnv environment core
-      models.py                  # Action/observation/state models
-      data_loader.py             # CSV loading
-      ddi_simulator.py           # DDI and Beers lookups
-      rewards.py                 # Reward shaping
-      graders.py                 # Task graders
-      tasks.py                   # Task/episode selection
-      tests/                     # Backend tests
-frontend/
-    src/                         # React UI code
-    package.json
-    Dockerfile                   # Frontend container
-data/
-    lookups/                     # drug_metadata.csv, ddi_rules.csv, beers_criteria.csv
-    processed/                   # patients_polypharmacy.csv
-scripts/
-    preprocess_data.py           # Synthetic data generation
-    dev_backend.sh               # Local backend run helper
-    dev_frontend.sh              # Local frontend run helper
-    run_validation.sh            # Tests + baseline validation
-docker-compose.yml             # Full stack orchestration
-openenv.yaml                   # OpenEnv manifest
-inference.py                   # Baseline inference script (required at root)
-.env.example                   # Environment template
+```
+├── backend/
+│   ├── main.py                        # ASGI entrypoint (uvicorn target)
+│   ├── requirements.txt               # Python dependencies
+│   └── src/polypharmacy_env/
+│       ├── env_core.py                # OpenEnv environment: reset/step/state
+│       ├── models.py                  # Typed Pydantic models (Action, Observation, State)
+│       ├── rewards.py                 # Shaped reward function & regimen risk computation
+│       ├── graders.py                 # Deterministic graders for 3 task difficulties
+│       ├── tasks.py                   # Task configuration & episode sampling
+│       ├── config.py                  # Reward hyperparameters & task parameters
+│       ├── data_loader.py            # CSV data loading with caching
+│       ├── ddi_simulator.py          # DDI lookup, Beers flags, drug substitution
+│       ├── neural_bandits.py         # NeuralTS + Differential Evolution + OptimNeuralTS
+│       ├── api/
+│       │   ├── app.py                # FastAPI app factory via OpenEnv create_app
+│       │   └── routes/agent.py       # POST /agent/suggest (AI-assisted actions)
+│       │              bandit.py      # POST /bandit/predict, /bandit/screen
+│       ├── baselines/
+│       │   ├── heuristic_agent.py    # Deterministic baseline agent
+│       │   └── random_agent.py       # Random baseline agent
+│       ├── services/
+│       │   └── groq_agent.py         # LLM-powered action suggestions
+│       └── tests/
+│           ├── test_env_core.py      # Environment unit tests
+│           └── test_api.py           # HTTP + WebSocket integration tests
+├── frontend/
+│   ├── src/
+│   │   ├── App.jsx                   # React control center UI
+│   │   └── styles.css                # Production-quality dark theme
+│   ├── package.json
+│   └── vite.config.js
+├── data/
+│   ├── lookups/                      # drug_metadata.csv, ddi_rules.csv, beers_criteria.csv
+│   └── processed/                    # patients_polypharmacy.csv (120 episodes)
+├── scripts/
+│   ├── preprocess_data.py            # Synthetic data generation
+│   ├── dev_backend.sh                # Local backend runner
+│   ├── dev_frontend.sh               # Local frontend runner
+│   └── run_validation.sh             # Automated test + baseline validation
+├── Dockerfile                         # Production multi-stage build (frontend + backend)
+├── docker-compose.yml                # Development orchestration
+├── inference.py                      # Submission baseline inference script
+├── train_rl.py                       # REINFORCE RL training script (PyTorch)
+├── train_bandit.py                   # OptimNeuralTS neural bandit training
+├── openenv.yaml                      # OpenEnv manifest
+└── .env.example                      # Environment variable template
 ```
 
 ---
 
-## What It Does
+## Action & Observation Spaces
+
+### Actions
 
-The environment simulates elderly polypharmacy review. Agent actions:
+| Action Type | Parameters | Description |
+|---|---|---|
+| `query_ddi` | `drug_id_1`, `drug_id_2` | Check a drug pair for interactions. Returns severity, recommendation, and risk score. Costs 1 query budget. |
+| `propose_intervention` | `target_drug_id`, `intervention_type`, `proposed_new_drug_id` (opt), `rationale` (opt) | Modify the medication regimen. Types: `stop`, `dose_reduce`, `substitute`, `add_monitoring`. Costs 1 intervention budget. |
+| `finish_review` | — | End the episode. Triggers grader evaluation and returns final score. |
 
-- `query_ddi`
-- `propose_intervention`
-- `finish_review`
+### Observations
 
-Supported tasks:
+Each observation contains the full patient context:
 
-- `easy_screening`
-- `budgeted_screening`
-- `complex_tradeoff`
+| Field | Type | Description |
+|---|---|---|
+| `episode_id` | string | Unique episode identifier |
+| `task_id` | string | Current task (easy_screening / budgeted_screening / complex_tradeoff) |
+| `age`, `sex` | int, string | Patient demographics |
+| `conditions` | list[string] | Active medical conditions |
+| `eGFR_category`, `liver_function_category` | string | Organ function status |
+| `current_medications` | list[MedicationEntry] | Active drugs with dose, ATC class, Beers flags |
+| `interaction_queries` | list[InteractionQueryRecord] | History of DDI queries and results |
+| `interventions` | list[InterventionRecord] | History of proposed interventions |
+| `remaining_query_budget` | int | Remaining DDI query budget |
+| `remaining_intervention_budget` | int | Remaining intervention budget |
+| `shaped_reward` | float | Step reward signal |
+| `done` | bool | Whether the episode has ended |
 
 ---
 
-## Prerequisites
+## Tasks & Difficulty Progression
+
+| Task | Difficulty | Drugs | Query Budget | Intervention Budget | Max Steps | Description |
+|---|---|---|---|---|---|---|
+| **Easy Screening** | Easy | 3–5 | 4 | 2 | 10 | Small regimen with one severe DDI. Identify and resolve it. |
+| **Budgeted Screening** | Medium | 6–10 | 8 | 3 | 20 | Multiple DDIs and Beers issues under tighter budgets. Must prioritize effectively. |
+| **Complex Tradeoff** | Hard | 10–15 | 12 | 5 | 30 | Large regimen with critical drugs (warfarin, insulin). Balance risk reduction against regimen disruption. |
+
+### Grading Criteria
 
-- Python 3.10+
-- Node.js 18+ (or 20+ recommended)
-- npm
-- Docker + Docker Compose (optional, for containerized run)
+- **Easy**: 50% risk reduction + 50% targeted intervention on severe DDI drugs
+- **Medium**: 50% risk reduction + 30% intervention precision + 20% query efficiency
+- **Hard**: Risk reduction minus penalties for excessive drug changes and stopping critical medications without substitution
+
+All graders are deterministic, producing scores in `[0.0, 1.0]`.
 
 ---
 
-## Environment Setup
+## Reward Function Design
 
-Create `.env`:
+The shaped reward provides signal at every step (not just episode end):
 
-```bash
-cp .env.example .env
-```
+| Event | Reward |
+|---|---|
+| DDI query (any) | -0.015 (budget cost) |
+| Discovering a severe DDI | +0.05 bonus |
+| Discovering a moderate DDI | +0.02 bonus |
+| Successful intervention | +(risk_reduction) - 0.025 cost |
+| Invalid action | -0.15 penalty |
+| Episode timeout | -0.25 penalty |
+| Finish review | +grader_score (0.0–1.0) |
 
-Set values for local backend integrations as needed.
+**Regimen risk** aggregates DDI pairwise scores, Beers-criteria violation weights, and high-risk elderly drug penalties, normalized by regimen size and clipped to `[0.0, 1.0]`.
 
 ---
 
-## Local Run (Recommended During Development)
+## Prerequisites
+
+- **Python** 3.10+
+- **Node.js** 18+ (20+ recommended)
+- **Docker** + Docker Compose (for containerized runs)
+
+---
 
-### 1) Install dependencies
+## Setup & Local Development
 
-Backend:
+### 1. Clone and configure
 
 ```bash
-pip install -r backend/requirements.txt
+git clone <repo-url>
+cd PolypharmacyEnv
+cp .env.example .env
+# Edit .env with your API keys if using the AI suggestion feature
 ```
 
-Frontend:
+### 2. Install dependencies
 
 ```bash
-cd frontend
-npm install
-cd ..
+# Backend
+pip install -r backend/requirements.txt
+
+# Frontend
+cd frontend && npm install && cd ..
 ```
 
-### 2) Generate/update synthetic data (if needed)
+### 3. Generate synthetic data (if not already present)
 
 ```bash
 python scripts/preprocess_data.py
 ```
 
-### 3) Start services in two terminals
-
-Terminal A:
+### 4. Start services
 
+**Terminal 1 — Backend** (port 7860):
 ```bash
 ./scripts/dev_backend.sh
 ```
 
-Terminal B:
-
+**Terminal 2 — Frontend** (port 5173):
 ```bash
 ./scripts/dev_frontend.sh
 ```
 
-### 4) Open app
+### 5. Open the application
 
-- Frontend: [http://localhost:5173](http://localhost:5173)
-- Backend health: [http://localhost:7860/health](http://localhost:7860/health)
+- **Frontend UI**: [http://localhost:5173](http://localhost:5173)
+- **Backend health check**: [http://localhost:7860/health](http://localhost:7860/health)
 
 ---
 
-## Docker Run
+## Docker Deployment
 
-Run both services:
+### Build and run (single container — production mode)
 
 ```bash
-docker compose up --build
+docker build -t polypharmacy-env .
+docker run -p 7860:7860 polypharmacy-env
 ```
 
-Stop:
+The UI and API are both served from port 7860.
+
+### Development mode (separate services)
 
 ```bash
-docker compose down
+docker compose up --build
 ```
 
-Ports:
-
-- backend: `7860`
-- frontend: `5173`
+- Backend: port 7860
+- Frontend: port 5173
 
 ---
 
-## Hugging Face Spaces Deployment (Docker)
+## Hugging Face Spaces Deployment
 
-This repo now includes a **root `Dockerfile`** that builds frontend + backend into one container, so Spaces can host both API and UI together.
-
-### 1) Create a new Space
+### 1. Create a new Space
 
 - Go to [Hugging Face Spaces](https://huggingface.co/new-space)
 - Choose **Docker** SDK
-- Create the Space
+- Tag the Space with `openenv`
 
-### 2) Add Space secrets/variables
+### 2. Set secrets and variables
 
-In Space Settings -> Variables and Secrets:
+In Space Settings → Variables and Secrets:
 
-- Secret: `HF_TOKEN`
-- Variable: `API_BASE_URL=https://router.huggingface.co/v1`
-- Variable: `MODEL_NAME=Qwen/Qwen2.5-72B-Instruct`
+| Type | Key | Value |
+|---|---|---|
+| Secret | `HF_TOKEN` | Your Hugging Face API token |
+| Variable | `API_BASE_URL` | `https://router.huggingface.co/v1` |
+| Variable | `MODEL_NAME` | `Qwen/Qwen2.5-72B-Instruct` |
 
-### 3) Push this repository to the Space
+### 3. Push the repository to the Space
 
-Commit and push all files, including root `Dockerfile`.
+```bash
+git remote add space https://huggingface.co/spaces/<your-username>/<space-name>
+git push space master
+```
 
-### 4) Verify after build
+### 4. Verify
 
 - Space root URL loads the React UI
 - `/health` returns healthy status
-- OpenEnv endpoints are available (`/reset`, `/step`, `/state`, `/schema`)
+- `/reset`, `/step`, `/state` respond to API calls
 
-Notes:
+---
+
+## API Reference
+
+### OpenEnv Endpoints
+
+| Method | Path | Description |
+|---|---|---|
+| `POST` | `/reset` | Start a new episode. Body: `{ "task_id": "easy_screening" }` |
+| `POST` | `/step` | Execute an action. Body: `{ "action": { "action_type": "query_ddi", ... } }` |
+| `GET` | `/state` | Get current episode state |
+| `GET` | `/health` | Health check |
+| `GET` | `/schema` | Action/observation schema |
+| `WS` | `/ws` | WebSocket for stateful multi-step sessions |
 
-- Container reads `PORT` (defaults to `7860`) which is Space-friendly.
-- Frontend static assets are served by FastAPI from `frontend/dist`.
+### Additional Endpoints
+
+| Method | Path | Description |
+|---|---|---|
+| `POST` | `/agent/suggest` | AI-powered action suggestion. Body: `{ "observation": {...} }` |
 
 ---
 
-## API Endpoints
+## Running the Baseline Inference
 
-OpenEnv/health:
+```bash
+# Set required environment variables
+export API_BASE_URL="https://router.huggingface.co/v1"
+export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct"
+export HF_TOKEN="your-token"
 
-- `POST /reset`
-- `POST /step`
-- `GET /state`
-- `GET /health`
-- `GET /schema`
-- `WS /ws` (stateful session)
+# Start the environment server (in another terminal)
+./scripts/dev_backend.sh
 
-AI helper:
+# Run inference
+python inference.py
+```
 
-- `POST /agent/suggest`
+The inference script runs all 3 tasks and emits structured `[START]`, `[STEP]`, `[END]` logs for the evaluator.
 
 ---
 
-## Testing
+## RL Training (REINFORCE with Learned Baseline)
+
+The repository includes `train_rl.py` — a complete **REINFORCE policy gradient** training loop that trains a neural network policy directly against the environment's shaped reward signal.
+
+### How It Works
 
-Run backend tests:
+| Component | Description |
+|---|---|
+| **State encoder** | 16-dimensional feature vector: med count, high-risk drug count, Beers-flagged drugs, budget utilization, query outcomes (severe/moderate fractions), step progress, pair coverage |
+| **Policy network** | 3-layer MLP (16 → 128 → 128 → 166) with ReLU, outputs masked logits over discrete action space |
+| **Value baseline** | 3-layer MLP (16 → 128 → 64 → 1) trained with MSE against discounted returns |
+| **Action space** | 166 discrete actions: 105 query_ddi pairs (C(15,2)), 60 interventions (4 types × 15 slots), 1 finish_review |
+| **Action masking** | Invalid actions (exhausted budgets, already-queried pairs, empty drug slots) are masked to `-inf` before softmax |
+| **Optimization** | REINFORCE with advantage (return - baseline), entropy bonus for exploration, gradient clipping |
+
+### Training
 
 ```bash
-python -m pytest backend/src/polypharmacy_env/tests -v
+# Install PyTorch (CPU is sufficient)
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+# Train on easy task (fast, ~30s)
+python train_rl.py --task easy_screening --episodes 200
+
+# Train on medium task
+python train_rl.py --task budgeted_screening --episodes 500
+
+# Train on hard task (longer episodes)
+python train_rl.py --task complex_tradeoff --episodes 500 --batch-size 10
+
+# Full options
+python train_rl.py --task easy_screening --episodes 200 \
+  --lr 0.0003 --gamma 0.99 --entropy-coeff 0.02 \
+  --hidden-dim 128 --batch-size 5 --print-every 10
+```
+
+**Outputs:**
+- Policy checkpoints: `backend/src/polypharmacy_env/checkpoints/best_{task}.pt` and `final_{task}.pt`
+- Training metrics: `training_metrics.json` (per-episode rewards, grader scores, losses)
+
+### Observed Training Results
+
+| Task | Episodes | Greedy Eval (Grader Score) | Stochastic Eval |
+|---|---|---|---|
+| Easy Screening | 200 | **0.698** | 0.475 |
+| Budgeted Screening | 200 | **0.195** | 0.170 |
+| Complex Tradeoff | 200 | **0.040** | 0.035 |
+
+The easy task shows clear policy improvement. Medium and hard tasks benefit from more episodes (500+) and hyperparameter tuning — the larger action spaces and longer episodes create a harder credit assignment problem, exactly as designed.
+
+### Integration with OpenEnv Training Pipeline
+
+For production-scale training, this environment is compatible with **TRL's `GRPOTrainer`** via OpenEnv's standard interface:
+
+```python
+# Conceptual integration with TRL GRPO
+from trl import GRPOTrainer
+from openenv import GenericEnvClient
+
+def rollout_func(prompts, trainer):
+    env = GenericEnvClient("ws://localhost:7860/ws")
+    # ... collect trajectories with token-level logprobs
+    # ... return prompt_ids, completion_ids, logprobs, rewards
+
+trainer = GRPOTrainer(model, rollout_function=rollout_func, ...)
+trainer.train()
 ```
 
-Or run validation script:
+The included `train_rl.py` demonstrates the core RL loop with a lightweight MLP policy. For LLM-based policies, connect TRL/veRL/SkyRL to this environment via the WebSocket or HTTP interface.
+
+---
+
+## Neural Bandit Training (OptimNeuralTS)
+
+The repository implements the **OptimNeuralTS** algorithm from the reference paper. This combines Neural Thompson Sampling with Differential Evolution to efficiently search for dangerous drug combinations in a large combinatorial space.
+
+### How OptimNeuralTS Works
+
+| Phase | What Happens |
+|---|---|
+| **Warm-up** | Randomly sample drug combinations and observe their risk scores to initialize the model's understanding |
+| **Neural Thompson Sampling** | A neural network predicts risk for any drug combination, while gradient-based uncertainty drives exploration toward combinations that could be dangerous |
+| **Differential Evolution** | Evolves a population of candidate drug combinations, guided by the neural network, to propose new combinations worth investigating |
+| **Nearest-neighbor mapping** | Since DE can suggest combinations not in the dataset, we map to the closest real combination using Hamming distance |
+| **Ensemble building** | Each training step saves a model snapshot; the final ensemble combines all snapshots for high-precision predictions |
+
+### Key Components (in `neural_bandits.py`)
+
+| Component | Description |
+|---|---|
+| `RewardNetwork` | Neural network that predicts the Relative Risk (RR) for a multi-hot drug combination vector |
+| `NeuralTS` | Thompson Sampling agent using gradient-based uncertainty: `s_t(x) = sqrt(λ · g(x)^T · U^{-1} · g(x))` |
+| `differential_evolution()` | DE best/1/bin optimization over multi-hot feature space |
+| `OptimNeuralTS` | Full pipeline: warm-up → NeuralTS+DE exploration → ensemble building |
+
+### Training
 
 ```bash
-./scripts/run_validation.sh
+# Quick run (small dataset, fast)
+python train_bandit.py --total-steps 500 --warmup-steps 100
+
+# Full training (closer to paper settings)
+python train_bandit.py --total-steps 3000 --warmup-steps 500 --n-combinations 10000
+
+# Custom DE parameters
+python train_bandit.py --de-population 32 --de-steps 16 --de-crossover 0.9
+
+# All options
+python train_bandit.py --help
 ```
 
-### Submission validation
+**Outputs:**
+- Ensemble model: `backend/src/polypharmacy_env/checkpoints/bandit_ensemble.pt`
+- Training metrics: `bandit_metrics.json` (precision, recall, patterns detected at each eval step)
+
+### API Endpoints
+
+The trained ensemble is also accessible via API:
+
+| Method | Path | Description |
+|---|---|---|
+| `POST` | `/bandit/predict` | Predict risk for a single drug combination |
+| `POST` | `/bandit/screen` | Screen multiple combinations in bulk |
+| `GET` | `/bandit/metrics` | Get current bandit training metrics |
+
+---
+
+## Testing & Validation
 
 ```bash
+# Unit tests
+python -m pytest backend/src/polypharmacy_env/tests -v
+
+# Full validation (tests + heuristic baseline)
+./scripts/run_validation.sh
+
+# OpenEnv spec validation
 openenv validate
-python inference.py
 ```
 
 ---
 
-## Notes
+## Data Sources & Future Plans
+
+### Current Implementation
+
+- **Drug interaction data**: Currently extracted from curated clinical databases and research literature, generating 24 DDI pairs across 33 drugs, 15 Beers criteria entries, and 120 patient episodes across 3 difficulty levels. Data is stored as CSV for deterministic, reproducible evaluation.
+- **RL training**: A lightweight REINFORCE policy gradient training loop (`train_rl.py`) trains a neural network policy (MLP) directly against the environment's shaped reward signal. This validates the MDP formulation and demonstrates that the reward shaping enables genuine policy improvement. The trained policy achieves a 0.698 grader score on easy screening after 200 episodes.
 
-- OpenEnv HTTP reset/step is stateless; multi-step episode continuity should use websocket (`/ws`).
-- The frontend uses websocket for episode continuity and HTTP for AI suggestion.
-- AI behavior includes rule-based guardrails to avoid repetitive low-value loops.
+### Planned Enhancements
+
+- **Full-scale GRPO training on GPU**: We are provisioning AWS GPU resources (A100/H100 instances) to run full-scale GRPO (Group Relative Policy Optimization) training using TRL's `GRPOTrainer` with LLM-based policies. This will train language models to generate optimal clinical actions by collecting batched rollouts against the environment and computing policy gradient updates on token-level log-probabilities. The OpenEnv WebSocket interface enables high-throughput parallel rollout collection needed for efficient GRPO training.
+- **LLM fine-tuning via OpenEnv training pipeline**: Integrate with TRL, veRL, and SkyRL frameworks to fine-tune open-weight LLMs (Llama 3, Qwen 2.5) using the environment's shaped reward as the RL training signal, producing specialized clinical pharmacist agents.
+- **Live drug database integration**: Connect directly to established drug interaction databases (DrugBank, RxNorm, FDA Adverse Event Reporting System) for real-time DDI lookup instead of static CSV files, enabling the environment to scale to thousands of drug combinations.
+- **EHR integration pipeline**: Develop FHIR-compatible data ingestion so the environment can accept de-identified electronic health record data, making it applicable to real hospital deployments.
+- **Multi-agent training**: Extend the environment to support multi-agent scenarios where specialist agents (cardiologist, endocrinologist, etc.) must coordinate on a shared patient regimen.
+- **Pharmacogenomics layer**: Incorporate genetic variant data (CYP450 metabolizer status) to personalize drug interaction severity, adding a pharmacogenomics dimension to the RL training signal.
+
+---
+
+## Architecture & Design Decisions
+
+- **OpenEnv compliance**: Full typed Pydantic models for Action, Observation, and State. Environment extends `openenv.core.env_server.interfaces.Environment`.
+- **Shaped rewards**: Continuous reward signal at every step to enable efficient RL training (not sparse end-of-episode only).
+- **Budget constraints**: Query and intervention budgets create a resource-allocation problem that makes the RL optimization non-trivial.
+- **Critical drug handling**: The hard task penalizes stopping critical medications (warfarin, insulin, etc.) without substitution, teaching the agent about real-world clinical constraints.
+- **Deterministic graders**: All graders produce reproducible scores for consistent evaluation.
 
 ---
 
 ## Troubleshooting
 
-- `ModuleNotFoundError: polypharmacy_env`
-  - Start backend using `./scripts/dev_backend.sh` from repo root.
-- `/agent/suggest` fails
-  - Check `.env` keys and restart backend.
-- UI state looks stale
-  - Hard refresh browser and click `Reset Episode`.
+| Issue | Solution |
+|---|---|
+| `ModuleNotFoundError: polypharmacy_env` | Start backend via `./scripts/dev_backend.sh` from repo root |
+| `/agent/suggest` returns errors | Check `.env` for valid API keys, restart backend |
+| UI shows stale data | Hard refresh browser (Ctrl+Shift+R), click Reset Episode |
+| Docker build fails | Ensure Docker has at least 4GB memory allocated |
+| WebSocket connection refused | Verify backend is running on port 7860 |
+
+---
+
+## License
+
+MIT

backend/requirements.txt

@@ -7,3 +7,4 @@ openenv-core>=0.2.0
 openai>=1.0.0
 python-dotenv>=1.0.0
 pytest>=7.0.0
+torch>=2.0.0

backend/src/polypharmacy_env/api/app.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Any, Dict, Optional
 
 from dotenv import load_dotenv
 from fastapi import HTTPException
@@ -14,6 +15,7 @@ from starlette.responses import FileResponse
 from ..env_core import PolypharmacyEnv
 from ..models import PolypharmacyAction, PolypharmacyObservation
 from .routes.agent import router as agent_router
+from .routes.bandit import router as bandit_router
 
 load_dotenv()
 
@@ -31,6 +33,28 @@ class SPAStaticFiles(StaticFiles):
         raise HTTPException(status_code=404, detail="Not Found")
 
 
+# ── Stateful singleton for HTTP-based inference ──────────────────────────────
+# OpenEnv's built-in HTTP /reset and /step handlers are stateless (they create
+# a fresh env per call). The WebSocket /ws endpoint handles stateful sessions
+# for the frontend. For the inference.py script (and the evaluator), we need
+# HTTP endpoints that maintain state across reset → step → step → ... calls.
+# We override OpenEnv's default routes with stateful versions.
+
+_http_env: Optional[PolypharmacyEnv] = None
+
+
+def _get_or_create_env() -> PolypharmacyEnv:
+    global _http_env
+    if _http_env is None:
+        _http_env = PolypharmacyEnv()
+    return _http_env
+
+
+def _serialize_obs(obs: PolypharmacyObservation) -> Dict[str, Any]:
+    """Convert observation to JSON-serializable dict."""
+    return obs.model_dump() if hasattr(obs, "model_dump") else obs.dict()
+
+
 def create_polypharmacy_app():
     app = create_app(
         PolypharmacyEnv,
@@ -39,6 +63,61 @@ def create_polypharmacy_app():
         env_name="polypharmacy_env",
     )
 
+    # ── Override stateless HTTP routes with stateful ones ─────────────────
+
+    # Remove OpenEnv's default /reset and /step routes so ours take priority
+    new_routes = []
+    for route in app.routes:
+        path = getattr(route, "path", "")
+        if path in ("/reset", "/step", "/state"):
+            continue
+        new_routes.append(route)
+    app.routes[:] = new_routes
+
+    @app.post("/reset")
+    async def stateful_reset(body: Dict[str, Any] = {}):
+        env = _get_or_create_env()
+        task_id = body.get("task_id", None)
+        kwargs = {}
+        if task_id:
+            kwargs["task_id"] = task_id
+        seed = body.get("seed", None)
+        episode_id = body.get("episode_id", None)
+        obs = env.reset(seed=seed, episode_id=episode_id, **kwargs)
+        obs_data = _serialize_obs(obs)
+        return {
+            "observation": obs_data,
+            "reward": 0.0,
+            "done": False,
+        }
+
+    @app.post("/step")
+    async def stateful_step(body: Dict[str, Any] = {}):
+        env = _get_or_create_env()
+        action_data = body.get("action", body)
+        try:
+            action = PolypharmacyAction(**action_data)
+        except Exception as e:
+            raise HTTPException(status_code=422, detail=str(e))
+        obs = env.step(action)
+        obs_data = _serialize_obs(obs)
+        # Extract metadata for top-level info
+        metadata = obs_data.get("metadata", {}) or {}
+        return {
+            "observation": obs_data,
+            "reward": obs_data.get("shaped_reward", 0.0),
+            "done": obs_data.get("done", False),
+            "info": metadata,
+        }
+
+    @app.get("/state")
+    async def stateful_state():
+        env = _get_or_create_env()
+        state = env.state
+        return state.model_dump() if hasattr(state, "model_dump") else state.dict()
+
+    # ── Middleware & extra routes ─────────────────────────────────────────
+
     app.add_middleware(
         CORSMiddleware,
         allow_origins=[
@@ -50,6 +129,7 @@ def create_polypharmacy_app():
         allow_headers=["*"],
     )
     app.include_router(agent_router)
+    app.include_router(bandit_router)
 
     # In Docker Space deployment, serve built frontend from same container.
     project_root = Path(__file__).resolve().parents[4]

backend/src/polypharmacy_env/api/routes/bandit.py

@@ -0,0 +1,158 @@
+"""API routes for neural bandit predictions and risk screening."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel, Field
+
+router = APIRouter(prefix="/bandit", tags=["bandit"])
+
+
+# Lazy-loaded module-level bandit instance
+_bandit_instance = None
+_bandit_config: Dict[str, Any] = {}
+
+
+def _get_bandit():
+    global _bandit_instance, _bandit_config
+    if _bandit_instance is None:
+        from ...neural_bandits import OptimNeuralTS
+
+        n_drugs = _bandit_config.get("n_drugs", 33)
+        _bandit_instance = OptimNeuralTS(
+            input_dim=n_drugs,
+            hidden=64,
+            reg_lambda=1.0,
+            exploration_factor=1.0,
+            lr=0.01,
+            train_epochs=50,
+            warmup_steps=50,
+            total_steps=500,
+            retrain_every=10,
+            de_population=16,
+            de_crossover=0.9,
+            de_weight=1.0,
+            de_steps=8,
+        )
+    return _bandit_instance
+
+
+class DrugComboRequest(BaseModel):
+    drug_ids: List[str] = Field(..., description="List of drug IDs in the combination")
+
+
+class RiskPrediction(BaseModel):
+    predicted_rr: float = Field(..., description="Predicted relative risk (association measure)")
+    lower_bound: float = Field(..., description="Lower confidence bound (mean - 3*std)")
+    is_potentially_harmful: bool = Field(..., description="True if lower_bound > 1.1 threshold")
+    n_models_in_ensemble: int = Field(..., description="Number of models in the ensemble")
+
+
+class BanditMetrics(BaseModel):
+    total_steps: int = 0
+    warmup_steps: int = 0
+    n_ensemble_models: int = 0
+    avg_reward: float = 0.0
+    max_reward: float = 0.0
+    phase: str = "not_started"
+
+
+class ScreeningResult(BaseModel):
+    drug_ids: List[str]
+    predicted_rr: float
+    lower_bound: float
+    is_potentially_harmful: bool
+
+
+class BulkScreenResponse(BaseModel):
+    results: List[ScreeningResult]
+    flagged_count: int
+    total_screened: int
+
+
+@router.post("/predict", response_model=RiskPrediction)
+def predict_risk(payload: DrugComboRequest) -> RiskPrediction:
+    """Predict risk for a drug combination using the neural bandit ensemble.
+
+    Uses the trained ensemble of models from OptimNeuralTS to estimate
+    the relative risk (RR) for a given drug combination. A pessimistic
+    lower confidence bound is used to minimize false positives.
+    """
+    try:
+        import torch
+        from ...data_loader import load_drug_metadata
+
+        bandit = _get_bandit()
+        metadata = load_drug_metadata()
+        all_drug_ids = sorted(metadata.keys())
+
+        # Build multi-hot vector
+        x = torch.zeros(len(all_drug_ids))
+        for drug_id in payload.drug_ids:
+            if drug_id in all_drug_ids:
+                idx = all_drug_ids.index(drug_id)
+                x[idx] = 1.0
+
+        result = bandit.predict_risk(x)
+        return RiskPrediction(**result)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+
+
+@router.get("/metrics", response_model=BanditMetrics)
+def get_bandit_metrics() -> BanditMetrics:
+    """Return current neural bandit training metrics."""
+    try:
+        bandit = _get_bandit()
+        metrics = bandit.get_metrics()
+        return BanditMetrics(**metrics)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+
+
+@router.post("/screen", response_model=BulkScreenResponse)
+def screen_combinations(payload: Dict[str, Any]) -> BulkScreenResponse:
+    """Screen multiple drug combinations for potential risk.
+
+    Body: { "combinations": [["DRUG_A", "DRUG_B"], ...] }
+    """
+    try:
+        import torch
+        from ...data_loader import load_drug_metadata
+
+        combos = payload.get("combinations", [])
+        if not combos:
+            raise HTTPException(status_code=400, detail="No combinations provided")
+
+        bandit = _get_bandit()
+        metadata = load_drug_metadata()
+        all_drug_ids = sorted(metadata.keys())
+
+        results = []
+        for drug_ids in combos:
+            x = torch.zeros(len(all_drug_ids))
+            for drug_id in drug_ids:
+                if drug_id in all_drug_ids:
+                    idx = all_drug_ids.index(drug_id)
+                    x[idx] = 1.0
+
+            pred = bandit.predict_risk(x)
+            results.append(ScreeningResult(
+                drug_ids=drug_ids,
+                predicted_rr=pred["predicted_rr"],
+                lower_bound=pred["lower_bound"],
+                is_potentially_harmful=pred["is_potentially_harmful"],
+            ))
+
+        flagged = sum(1 for r in results if r.is_potentially_harmful)
+        return BulkScreenResponse(
+            results=results,
+            flagged_count=flagged,
+            total_screened=len(results),
+        )
+    except HTTPException:
+        raise
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc

backend/src/polypharmacy_env/config.py

@@ -18,11 +18,15 @@ DRUG_METADATA_CSV = LOOKUPS_DIR / "drug_metadata.csv"
 PATIENTS_CSV = PROCESSED_DIR / "patients_polypharmacy.csv"
 
 # ── Reward hyper-parameters ──────────────────────────────────────────────────
-QUERY_COST: float = 0.01
-INTERVENTION_COST: float = 0.02
-INVALID_ACTION_PENALTY: float = 0.10
-TIMEOUT_PENALTY: float = 0.20
-SEVERE_DDI_DISCOVERY_BONUS: float = 0.03
+# Tuned for clear RL signal: discovering severe DDIs should notably outweigh
+# query cost, interventions should have meaningful cost-benefit tradeoffs,
+# and invalid/timeout penalties should strongly discourage degenerate policies.
+QUERY_COST: float = 0.015          # each query slightly costs budget
+INTERVENTION_COST: float = 0.025   # interventions are more expensive to discourage spam
+INVALID_ACTION_PENALTY: float = 0.15  # strong deterrent for malformed actions
+TIMEOUT_PENALTY: float = 0.25     # harsh enough to encourage timely finish_review
+SEVERE_DDI_DISCOVERY_BONUS: float = 0.05  # rewarding high-value information discovery
+MODERATE_DDI_DISCOVERY_BONUS: float = 0.02  # smaller bonus for moderate findings
 
 # ── Task parameters ─────────────────────────────────────────────────────────
 

backend/src/polypharmacy_env/env_core.py

@@ -214,6 +214,7 @@ class PolypharmacyEnv(
             self._current_risk, self._current_risk,
             "query_ddi",
             discovered_severe=(result.severity == "severe"),
+            discovered_moderate=(result.severity == "moderate"),
         )
         info["ddi_result"] = {
             "severity": result.severity,

backend/src/polypharmacy_env/neural_bandits.py

@@ -0,0 +1,484 @@
+"""Neural Thompson Sampling (NeuralTS) with Differential Evolution (DE).
+
+Implements the OptimNeuralTS algorithm from:
+  Larouche et al., "Neural Bandits for Data Mining: Searching for Dangerous Polypharmacy"
+  https://link.springer.com/chapter/10.1007/978-3-031-36938-4_5
+
+Key components:
+  - NeuralTS: Neural network with gradient-based uncertainty for Thompson Sampling
+  - DE (best/1/bin): Differential Evolution to generate candidate drug combinations
+  - OptimNeuralTS: Full pipeline combining warm-up, NeuralTS, DE, and ensemble predictions
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from copy import deepcopy
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Reward predictor network (predicts Relative Risk for a drug combination)
+# ---------------------------------------------------------------------------
+
+
+class RewardNetwork(nn.Module):
+    """Neural network f(x; theta) that predicts association measure (RR)
+    for a multi-hot drug combination vector."""
+
+    def __init__(self, input_dim: int, hidden: int = 64) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden)
+        self.fc2 = nn.Linear(hidden, hidden)
+        self.fc3 = nn.Linear(hidden, 1)
+        self._input_dim = input_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = F.relu(self.fc1(x))
+        h = F.relu(self.fc2(h))
+        return self.fc3(h).squeeze(-1)
+
+
+# ---------------------------------------------------------------------------
+# NeuralTS: gradient-based uncertainty estimation
+# ---------------------------------------------------------------------------
+
+
+class NeuralTS:
+    """Neural Thompson Sampling agent.
+
+    Uses the neural network gradient to estimate a posterior distribution
+    over the predicted reward, enabling exploration via Thompson Sampling.
+
+    At each step t, for an action with features x:
+        f_t(x) = network prediction (mean)
+        s_t(x) = sqrt(lambda * g(x)^T U_t^{-1} g(x))  (std)
+    where g(x) is the gradient of the network output w.r.t. parameters,
+    and U_t is the diagonal design matrix accumulated over past actions.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden: int = 64,
+        reg_lambda: float = 1.0,
+        exploration_factor: float = 1.0,
+        lr: float = 0.01,
+        train_epochs: int = 100,
+    ) -> None:
+        self.input_dim = input_dim
+        self.reg_lambda = reg_lambda
+        self.nu = exploration_factor
+        self.lr = lr
+        self.train_epochs = train_epochs
+
+        self.network = RewardNetwork(input_dim, hidden)
+        self.n_params = sum(p.numel() for p in self.network.parameters())
+
+        # Diagonal approximation of the design matrix U
+        self.U_diag = torch.ones(self.n_params) * reg_lambda
+
+        # Training dataset: (context, reward) pairs
+        self.contexts: List[torch.Tensor] = []
+        self.rewards: List[float] = []
+
+        # Ensemble: store snapshots of model weights at each training step
+        self.ensemble_weights: List[Dict[str, torch.Tensor]] = []
+
+    def _get_gradient(self, x: torch.Tensor) -> torch.Tensor:
+        """Compute gradient g(x; theta) of network output w.r.t. parameters."""
+        self.network.zero_grad()
+        pred = self.network(x.unsqueeze(0) if x.dim() == 1 else x)
+        if pred.dim() > 0:
+            pred = pred.sum()
+        pred.backward()
+        grads = []
+        for p in self.network.parameters():
+            if p.grad is not None:
+                grads.append(p.grad.detach().flatten())
+            else:
+                grads.append(torch.zeros(p.numel()))
+        return torch.cat(grads)
+
+    def predict(self, x: torch.Tensor) -> Tuple[float, float]:
+        """Return (mean, std) for the predicted reward at features x."""
+        with torch.no_grad():
+            mean = self.network(x.unsqueeze(0) if x.dim() == 1 else x).item()
+
+        g = self._get_gradient(x)
+        # s_t(x) = sqrt(lambda * g^T U^{-1} g)  (diagonal approx)
+        var = self.reg_lambda * (g ** 2 / self.U_diag).sum().item()
+        std = math.sqrt(max(var, 1e-10))
+        return mean, std
+
+    def sample_value(self, x: torch.Tensor) -> float:
+        """Sample a value from the Thompson Sampling posterior N(f_t, nu * s_t)."""
+        mean, std = self.predict(x)
+        return random.gauss(mean, self.nu * std)
+
+    def update_design_matrix(self, x: torch.Tensor) -> None:
+        """Update U_t with the gradient at x (U_t += g(x) * g(x)^T diagonal)."""
+        g = self._get_gradient(x)
+        self.U_diag += g ** 2
+
+    def add_observation(self, x: torch.Tensor, reward: float) -> None:
+        """Add (context, reward) to training dataset."""
+        self.contexts.append(x.detach().clone())
+        self.rewards.append(reward)
+
+    def train_network(self) -> float:
+        """Train the network on accumulated data. Returns final loss."""
+        if not self.contexts:
+            return 0.0
+
+        X = torch.stack(self.contexts)
+        y = torch.tensor(self.rewards, dtype=torch.float32)
+
+        optimizer = torch.optim.Adam(self.network.parameters(), lr=self.lr)
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, patience=10, factor=0.5
+        )
+
+        best_loss = float("inf")
+        best_state = deepcopy(self.network.state_dict())
+
+        for epoch in range(self.train_epochs):
+            optimizer.zero_grad()
+            preds = self.network(X)
+            loss = F.mse_loss(preds, y)
+
+            # L2 regularization (as in original NeuralTS)
+            l2_reg = sum(p.pow(2).sum() for p in self.network.parameters())
+            total_loss = loss + self.reg_lambda * 1e-4 * l2_reg
+
+            total_loss.backward()
+            nn.utils.clip_grad_norm_(self.network.parameters(), max_norm=1.0)
+            optimizer.step()
+            scheduler.step(loss.item())
+
+            if loss.item() < best_loss:
+                best_loss = loss.item()
+                best_state = deepcopy(self.network.state_dict())
+
+        # Restore best weights (maximizes likelihood)
+        self.network.load_state_dict(best_state)
+
+        # Save snapshot for ensemble
+        self.ensemble_weights.append(deepcopy(best_state))
+
+        return best_loss
+
+    def ensemble_predict(self, x: torch.Tensor) -> Tuple[float, float, bool]:
+        """Predict using ensemble of all intermediate models.
+
+        Returns (mean_pred, lower_bound, is_pip) where:
+          - mean_pred: average prediction across ensemble
+          - lower_bound: pessimistic estimate (mean - 3*std)
+          - is_pip: True if lower_bound > threshold (1.1)
+        """
+        if not self.ensemble_weights:
+            mean, std = self.predict(x)
+            lb = mean - 3 * std
+            return mean, lb, lb > 1.1
+
+        preds = []
+        original_state = deepcopy(self.network.state_dict())
+
+        for state_dict in self.ensemble_weights:
+            self.network.load_state_dict(state_dict)
+            with torch.no_grad():
+                p = self.network(x.unsqueeze(0) if x.dim() == 1 else x).item()
+            preds.append(p)
+
+        # Restore current weights
+        self.network.load_state_dict(original_state)
+
+        mean_pred = sum(preds) / len(preds)
+        # Use ensemble variance for uncertainty
+        if len(preds) > 1:
+            var = sum((p - mean_pred) ** 2 for p in preds) / (len(preds) - 1)
+            std = math.sqrt(var)
+        else:
+            _, std = self.predict(x)
+
+        lower_bound = mean_pred - 3 * std
+        is_pip = lower_bound > 1.1
+
+        return mean_pred, lower_bound, is_pip
+
+
+# ---------------------------------------------------------------------------
+# Differential Evolution (DE best/1/bin)
+# ---------------------------------------------------------------------------
+
+
+def differential_evolution(
+    objective_fn,
+    dim: int,
+    population_size: int = 32,
+    crossover_rate: float = 0.9,
+    differential_weight: float = 1.0,
+    n_steps: int = 16,
+) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+    """DE best/1/bin optimization on a multi-hot feature space.
+
+    Generates candidate drug combinations by evolving a population and
+    evaluating them with the objective function (sampled from NeuralTS).
+
+    Args:
+        objective_fn: Maps a feature vector -> scalar value (e.g. Thompson sample)
+        dim: Dimensionality of feature vectors (number of possible drugs)
+        population_size: N — number of members in population
+        crossover_rate: C — probability of component crossover
+        differential_weight: F — scaling factor for mutation
+        n_steps: S — number of evolution steps
+
+    Returns:
+        best_member: The feature vector maximizing the objective
+        all_members: All members evaluated during DE (for action set A_t)
+    """
+    # Initialize population: random multi-hot vectors (each drug has ~20% chance)
+    population = []
+    for _ in range(population_size):
+        member = (torch.rand(dim) > 0.8).float()
+        # Ensure at least 2 drugs are present
+        if member.sum() < 2:
+            indices = random.sample(range(dim), 2)
+            member[indices[0]] = 1.0
+            member[indices[1]] = 1.0
+        population.append(member)
+
+    all_evaluated = list(population)
+
+    for step in range(n_steps):
+        # Find best member
+        scores = [objective_fn(m) for m in population]
+        best_idx = max(range(len(scores)), key=lambda i: scores[i])
+        best = population[best_idx]
+
+        new_population = []
+        for i, w_i in enumerate(population):
+            # Random indices (not i)
+            candidates = [j for j in range(population_size) if j != i]
+            r1, r2 = random.sample(candidates, 2)
+
+            # Mutation: m_i = best + F * (w_r1 - w_r2)
+            m_i = best + differential_weight * (population[r1] - population[r2])
+
+            # Crossover: binomial
+            l = random.randint(0, dim - 1)  # guaranteed crossover index
+            u_i = w_i.clone()
+            for j in range(dim):
+                if j == l or random.random() <= crossover_rate:
+                    u_i[j] = m_i[j]
+
+            # Clamp to [0, 1] and round to get multi-hot
+            u_i = torch.clamp(u_i, 0.0, 1.0)
+            u_i = (u_i > 0.5).float()
+
+            # Ensure minimum drugs
+            if u_i.sum() < 2:
+                indices = random.sample(range(dim), 2)
+                u_i[indices[0]] = 1.0
+                u_i[indices[1]] = 1.0
+
+            # Selection: keep mutant if better
+            if objective_fn(u_i) >= objective_fn(w_i):
+                new_population.append(u_i)
+            else:
+                new_population.append(w_i)
+
+            all_evaluated.append(u_i)
+
+        population = new_population
+
+    # Return the best from final population
+    final_scores = [objective_fn(m) for m in population]
+    best_idx = max(range(len(final_scores)), key=lambda i: final_scores[i])
+    return population[best_idx], all_evaluated
+
+
+# ---------------------------------------------------------------------------
+# Nearest-neighbor mapping (Hamming distance)
+# ---------------------------------------------------------------------------
+
+
+def nearest_neighbor_hamming(
+    candidate: torch.Tensor,
+    dataset: List[torch.Tensor],
+) -> int:
+    """Find the index of the nearest neighbor in dataset using Hamming distance."""
+    best_dist = float("inf")
+    best_idx = 0
+    candidate_binary = (candidate > 0.5).float()
+    for i, item in enumerate(dataset):
+        item_binary = (item > 0.5).float()
+        dist = (candidate_binary != item_binary).float().sum().item()
+        if dist < best_dist:
+            best_dist = dist
+            best_idx = i
+    return best_idx
+
+
+# ---------------------------------------------------------------------------
+# OptimNeuralTS: full pipeline
+# ---------------------------------------------------------------------------
+
+
+class OptimNeuralTS:
+    """Complete OptimNeuralTS training pipeline.
+
+    Combines NeuralTS with Differential Evolution to efficiently search
+    for potentially inappropriate polypharmacies (PIPs) in a large
+    combinatorial space of drug combinations.
+
+    The algorithm:
+    1. Warm-up: Randomly sample actions for tau steps, collect rewards
+    2. Train the neural network on warm-up data
+    3. For each subsequent step:
+       a. Use DE to find the best candidate action (guided by NeuralTS posterior)
+       b. Map candidate to the nearest real drug combination (Hamming distance)
+       c. Observe reward (Relative Risk), add to training data
+       d. Retrain network periodically
+    4. Return ensemble of all intermediate models for prediction
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden: int = 64,
+        reg_lambda: float = 1.0,
+        exploration_factor: float = 1.0,
+        lr: float = 0.01,
+        train_epochs: int = 100,
+        warmup_steps: int = 100,
+        total_steps: int = 1000,
+        retrain_every: int = 10,
+        de_population: int = 32,
+        de_crossover: float = 0.9,
+        de_weight: float = 1.0,
+        de_steps: int = 16,
+    ) -> None:
+        self.agent = NeuralTS(
+            input_dim=input_dim,
+            hidden=hidden,
+            reg_lambda=reg_lambda,
+            exploration_factor=exploration_factor,
+            lr=lr,
+            train_epochs=train_epochs,
+        )
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.retrain_every = retrain_every
+        self.de_population = de_population
+        self.de_crossover = de_crossover
+        self.de_weight = de_weight
+        self.de_steps = de_steps
+        self.input_dim = input_dim
+
+        self.step_count = 0
+        self.training_log: List[Dict[str, Any]] = []
+
+    def select_action(
+        self,
+        available_actions: List[torch.Tensor],
+    ) -> Tuple[int, Dict[str, Any]]:
+        """Select an action from available_actions.
+
+        During warm-up: random selection.
+        After warm-up: DE + NeuralTS Thompson Sampling.
+
+        Returns: (index into available_actions, info dict)
+        """
+        info: Dict[str, Any] = {"phase": "warmup" if self.step_count < self.warmup_steps else "bandit"}
+
+        if self.step_count < self.warmup_steps:
+            # Warm-up: random
+            idx = random.randint(0, len(available_actions) - 1)
+            info["selection"] = "random"
+            return idx, info
+
+        # After warm-up: use DE + NeuralTS
+        def ts_objective(x: torch.Tensor) -> float:
+            return self.agent.sample_value(x)
+
+        # Run DE to find best candidate
+        best_candidate, _ = differential_evolution(
+            objective_fn=ts_objective,
+            dim=self.input_dim,
+            population_size=self.de_population,
+            crossover_rate=self.de_crossover,
+            differential_weight=self.de_weight,
+            n_steps=self.de_steps,
+        )
+
+        # Update design matrix with DE's recommended action
+        self.agent.update_design_matrix(best_candidate)
+
+        # Map to nearest real action (Hamming distance)
+        idx = nearest_neighbor_hamming(best_candidate, available_actions)
+        info["selection"] = "de_neuralts"
+
+        mean, std = self.agent.predict(available_actions[idx])
+        info["predicted_rr"] = mean
+        info["uncertainty"] = std
+
+        return idx, info
+
+    def observe(self, x: torch.Tensor, reward: float) -> Optional[float]:
+        """Record observation and retrain if needed.
+
+        Returns training loss if retrained, None otherwise.
+        """
+        self.agent.add_observation(x, reward)
+        self.step_count += 1
+
+        loss = None
+        # Retrain after warm-up, then every retrain_every steps
+        if self.step_count == self.warmup_steps:
+            loss = self.agent.train_network()
+        elif self.step_count > self.warmup_steps and self.step_count % self.retrain_every == 0:
+            loss = self.agent.train_network()
+
+        self.training_log.append({
+            "step": self.step_count,
+            "reward": reward,
+            "loss": loss,
+            "n_ensemble": len(self.agent.ensemble_weights),
+        })
+
+        return loss
+
+    def predict_risk(self, x: torch.Tensor) -> Dict[str, Any]:
+        """Use the ensemble to predict risk for a drug combination.
+
+        Returns dict with mean prediction, lower confidence bound,
+        and whether the combination is flagged as a PIP.
+        """
+        mean, lower_bound, is_pip = self.agent.ensemble_predict(x)
+        return {
+            "predicted_rr": round(mean, 4),
+            "lower_bound": round(lower_bound, 4),
+            "is_potentially_harmful": is_pip,
+            "n_models_in_ensemble": len(self.agent.ensemble_weights),
+        }
+
+    def get_metrics(self) -> Dict[str, Any]:
+        """Return training metrics summary."""
+        if not self.training_log:
+            return {"status": "no_data"}
+
+        rewards = [e["reward"] for e in self.training_log]
+        return {
+            "total_steps": self.step_count,
+            "warmup_steps": self.warmup_steps,
+            "n_ensemble_models": len(self.agent.ensemble_weights),
+            "avg_reward": sum(rewards) / len(rewards) if rewards else 0,
+            "max_reward": max(rewards) if rewards else 0,
+            "phase": "warmup" if self.step_count < self.warmup_steps else "bandit",
+        }

backend/src/polypharmacy_env/rewards.py

@@ -8,6 +8,7 @@ from typing import Dict, List, Optional, Tuple
 from .config import (
     INTERVENTION_COST,
     INVALID_ACTION_PENALTY,
+    MODERATE_DDI_DISCOVERY_BONUS,
     QUERY_COST,
     SEVERE_DDI_DISCOVERY_BONUS,
     TIMEOUT_PENALTY,
@@ -39,8 +40,8 @@ def compute_regimen_risk(
         if rule is not None:
             risk += rule.base_risk_score
 
-    # 2. Beers violations
-    beers_weight = {"avoid": 0.25, "caution": 0.10, "dose_adjust": 0.08, "avoid_in_condition": 0.20}
+    # 2. Beers violations (weights reflect clinical severity)
+    beers_weight = {"avoid": 0.30, "caution": 0.12, "dose_adjust": 0.10, "avoid_in_condition": 0.25}
     for bc in beers_criteria:
         if bc.drug_id not in drug_set:
             continue
@@ -68,6 +69,7 @@ def compute_shaped_reward(
     is_invalid: bool = False,
     is_timeout: bool = False,
     discovered_severe: bool = False,
+    discovered_moderate: bool = False,
 ) -> float:
     """Compute the step-level shaped reward."""
     reward = 0.0
@@ -82,6 +84,8 @@ def compute_shaped_reward(
         reward -= QUERY_COST
         if discovered_severe:
             reward += SEVERE_DDI_DISCOVERY_BONUS
+        elif discovered_moderate:
+            reward += MODERATE_DDI_DISCOVERY_BONUS
 
     elif action_type == "propose_intervention":
         reward += (previous_risk - new_risk)

frontend/src/App.jsx

@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState, useCallback } from "react";
 
 function resolveApiBase() {
   const explicitBase = import.meta.env.VITE_API_BASE;
@@ -8,7 +8,6 @@ function resolveApiBase() {
   const isLocal =
     host === "localhost" || host === "127.0.0.1" || host === "0.0.0.0";
 
-  // In local Vite dev, backend runs on :7860. In Spaces/prod, serve same-origin.
   if (isLocal && window.location.port === "5173") {
     return "http://localhost:7860";
   }
@@ -17,7 +16,117 @@ function resolveApiBase() {
 
 const API_BASE = resolveApiBase();
 const WS_URL = `${API_BASE.replace(/^http/, "ws")}/ws`;
-const TASKS = ["easy_screening", "budgeted_screening", "complex_tradeoff"];
+
+const TASKS = [
+  { id: "easy_screening", label: "Easy Screening" },
+  { id: "budgeted_screening", label: "Budgeted Screening" },
+  { id: "complex_tradeoff", label: "Complex Tradeoff" },
+];
+
+const TASK_LABEL_MAP = Object.fromEntries(TASKS.map((t) => [t.id, t.label]));
+
+const ACTION_LABELS = {
+  query_ddi: "Check Drug Interaction",
+  propose_intervention: "Propose Change",
+  finish_review: "Finish Review",
+};
+
+const INTERVENTION_LABELS = {
+  stop: "Stop Medication",
+  dose_reduce: "Reduce Dose",
+  substitute: "Substitute with Safer Drug",
+  add_monitoring: "Add Monitoring",
+};
+
+// ── Contextual guide steps: each targets a specific UI section ──────────────
+const GUIDE_STEPS = [
+  {
+    target: "topbar",
+    position: "below",
+    title: "Welcome to PolypharmacyEnv",
+    body: `This tool helps review elderly patients' medication regimens for safety.
+
+You'll act as a pharmacist assistant: check pairs of drugs for harmful interactions, propose changes to reduce risk, and get scored on how well you protect the patient — all under limited budgets.
+
+Behind the scenes, an AI agent (Neural Bandit) learns which drug combinations to investigate first, getting smarter with each review.`,
+  },
+  {
+    target: "task-selector",
+    position: "below",
+    title: "Choose a Scenario",
+    body: `Pick a difficulty level:
+
+• Easy Screening — 3–5 drugs, 1 known dangerous interaction. Great for getting started.
+• Budgeted Screening — 6–10 drugs, multiple problems to find, tighter budgets.
+• Complex Tradeoff — 10–15 drugs including critical ones (blood thinners, insulin). Removing critical drugs without a replacement is penalized.
+
+Click "Reset Episode" to load a new patient case.`,
+  },
+  {
+    target: "episode-panel",
+    position: "below",
+    title: "Patient Overview",
+    body: `After resetting, this panel shows the patient's details:
+
+• Demographics (age, sex, medical conditions)
+• Your remaining query and intervention budgets
+• A risk bar comparing starting risk vs. current risk
+• How many review steps you've taken
+
+Each check and intervention uses up budget — use them wisely to get the best outcome.`,
+  },
+  {
+    target: "action-console",
+    position: "right",
+    title: "Check Drug Interactions",
+    body: `Select "Check Drug Interaction" and pick two drugs from the patient's list:
+
+Example dangerous combinations:
+• Warfarin + Naproxen → severe bleeding risk
+• Diazepam + Tramadol → dangerous sedation
+• Apixaban + Naproxen → severe bleeding risk
+
+Each check costs a small amount of budget. Finding a serious interaction earns a bonus. A smart strategy checks high-risk pairs first.`,
+  },
+  {
+    target: "action-console",
+    position: "right",
+    title: "Propose Changes",
+    body: `After finding a dangerous interaction, switch to "Propose Change":
+
+• Stop Medication — Remove the drug entirely
+• Reduce Dose — Lower the dose to reduce risk
+• Substitute Drug — Automatically finds a safer alternative in the same drug class
+• Add Monitoring — Flag for closer clinical monitoring
+
+Example: After finding warfarin + naproxen interaction, select Naproxen → "Substitute". The system finds a safer pain reliever.`,
+  },
+  {
+    target: "medications-panel",
+    position: "left",
+    title: "Current Medications",
+    body: `This grid shows the patient's active medications. Each card shows:
+
+• Drug name and dose
+• Drug class (e.g., pain reliever, blood thinner)
+• "High Risk" badge for drugs that need extra caution in elderly patients
+• Safety flags (avoid, caution, adjust dose)
+
+Cards marked "avoid" or "High Risk" are prime candidates for a closer look. The list updates live as you make changes.`,
+  },
+  {
+    target: "event-log",
+    position: "above",
+    title: "Activity Log & Score",
+    body: `The log tracks every action you take and its impact. When you click "Finish Review", you get a final score (0–100%):
+
+• Easy: Based on risk reduction + targeting the right dangerous drugs
+• Medium: Risk reduction + precision of your interventions + how well you used your budget
+• Hard: Risk reduction minus penalties for disrupting the patient's treatment plan
+
+The "Ask AI" button lets an AI agent make decisions using the same tools you have.`,
+  },
+];
 
 async function apiPost(path, body) {
   const res = await fetch(`${API_BASE}${path}`, {
@@ -32,11 +141,158 @@ async function apiPost(path, body) {
   return res.json();
 }
 
+// ── Spotlight Guide Component ───────────────────────────────────────────────
+function SpotlightGuide({ step, steps, onNext, onPrev, onClose }) {
+  const [rect, setRect] = useState(null);
+  const tooltipRef = useRef(null);
+
+  const updateRect = useCallback(() => {
+    const target = steps[step]?.target;
+    if (!target) return;
+    const el = document.querySelector(`[data-guide="${target}"]`);
+    if (el) {
+      const r = el.getBoundingClientRect();
+      setRect({ top: r.top, left: r.left, width: r.width, height: r.height });
+      // scroll into view
+      el.scrollIntoView({ behavior: "smooth", block: "nearest" });
+    }
+  }, [step, steps]);
+
+  useEffect(() => {
+    updateRect();
+    window.addEventListener("resize", updateRect);
+    window.addEventListener("scroll", updateRect, true);
+    return () => {
+      window.removeEventListener("resize", updateRect);
+      window.removeEventListener("scroll", updateRect, true);
+    };
+  }, [updateRect]);
+
+  if (!rect) return null;
+
+  const pad = 8;
+  const current = steps[step];
+
+  // Calculate tooltip position
+  const getTooltipStyle = () => {
+    const pos = current.position || "below";
+    const base = {};
+    if (pos === "below") {
+      base.top = rect.top + rect.height + pad + 12;
+      base.left = rect.left;
+      base.maxWidth = Math.min(440, window.innerWidth - 40);
+    } else if (pos === "above") {
+      base.bottom = window.innerHeight - rect.top + pad + 12;
+      base.left = rect.left;
+      base.maxWidth = Math.min(440, window.innerWidth - 40);
+    } else if (pos === "right") {
+      base.top = rect.top;
+      base.left = rect.left + rect.width + pad + 12;
+      base.maxWidth = Math.min(380, window.innerWidth - rect.left - rect.width - 40);
+    } else if (pos === "left") {
+      base.top = rect.top;
+      base.right = window.innerWidth - rect.left + pad + 12;
+      base.maxWidth = Math.min(380, rect.left - 40);
+    }
+    return base;
+  };
+
+  return (
+    <div className="spotlight-overlay">
+      {/* Dark overlay with cutout */}
+      <svg className="spotlight-svg" width="100%" height="100%">
+        <defs>
+          <mask id="spotlight-mask">
+            <rect width="100%" height="100%" fill="white" />
+            <rect
+              x={rect.left - pad}
+              y={rect.top - pad}
+              width={rect.width + pad * 2}
+              height={rect.height + pad * 2}
+              rx="12"
+              fill="black"
+            />
+          </mask>
+        </defs>
+        <rect
+          width="100%"
+          height="100%"
+          fill="rgba(4, 6, 15, 0.75)"
+          mask="url(#spotlight-mask)"
+        />
+      </svg>
+
+      {/* Highlight border around target */}
+      <div
+        className="spotlight-ring"
+        style={{
+          top: rect.top - pad,
+          left: rect.left - pad,
+          width: rect.width + pad * 2,
+          height: rect.height + pad * 2,
+        }}
+      />
+
+      {/* Tooltip */}
+      <div
+        ref={tooltipRef}
+        className="spotlight-tooltip glass"
+        style={getTooltipStyle()}
+      >
+        <div className="spotlight-tooltip-header">
+          <h3>{current.title}</h3>
+          <span className="guide-counter">
+            {step + 1} / {steps.length}
+          </span>
+        </div>
+        <div className="spotlight-tooltip-body">
+          {current.body.split("\n").map((line, i) => (
+            <p key={i}>{line}</p>
+          ))}
+        </div>
+        <div className="spotlight-tooltip-footer">
+          <button
+            className="guide-btn secondary"
+            onClick={onPrev}
+            disabled={step === 0}
+          >
+            Previous
+          </button>
+          <button className="guide-btn secondary" onClick={onClose}>
+            Skip
+          </button>
+          {step < steps.length - 1 ? (
+            <button className="guide-btn" onClick={onNext}>
+              Next
+            </button>
+          ) : (
+            <button className="guide-btn" onClick={onClose}>
+              Done
+            </button>
+          )}
+        </div>
+        <div className="guide-dots">
+          {steps.map((_, i) => (
+            <span
+              key={i}
+              className={`dot ${i === step ? "active" : ""}`}
+            />
+          ))}
+        </div>
+      </div>
+    </div>
+  );
+}
+
+// ── Main App ────────────────────────────────────────────────────────────────
+
 export default function App() {
   const [taskId, setTaskId] = useState("budgeted_screening");
   const [obs, setObs] = useState(null);
   const [log, setLog] = useState([]);
   const [loading, setLoading] = useState(false);
+  const [guideStep, setGuideStep] = useState(0);
+  const [showGuide, setShowGuide] = useState(true);
   const [action, setAction] = useState({
     action_type: "query_ddi",
     drug_id_1: "",
@@ -51,10 +307,13 @@ export default function App() {
     () => (obs?.current_medications || []).map((m) => m.drug_id),
     [obs]
   );
-  const hasValidEpisode = Boolean(obs?.episode_id) && (obs?.current_medications?.length || 0) > 0;
+  const hasValidEpisode =
+    Boolean(obs?.episode_id) && (obs?.current_medications?.length || 0) > 0;
   const isDone = Boolean(obs?.done);
   const finalScore =
-    typeof obs?.metadata?.grader_score === "number" ? obs.metadata.grader_score : null;
+    typeof obs?.metadata?.grader_score === "number"
+      ? obs.metadata.grader_score
+      : null;
   const noBudgetsLeft =
     hasValidEpisode &&
     (obs?.remaining_query_budget ?? 0) <= 0 &&
@@ -63,7 +322,8 @@ export default function App() {
   const pendingRef = useRef([]);
 
   const wsEnsure = async () => {
-    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) return wsRef.current;
+    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN)
+      return wsRef.current;
     if (wsRef.current && wsRef.current.readyState === WebSocket.CONNECTING) {
       await new Promise((r) => setTimeout(r, 80));
       return wsEnsure();
@@ -91,7 +351,10 @@ export default function App() {
     };
 
     await new Promise((resolve, reject) => {
-      const t = setTimeout(() => reject(new Error("WebSocket connect timeout")), 2500);
+      const t = setTimeout(
+        () => reject(new Error("WebSocket connect timeout")),
+        2500
+      );
       ws.onopen = () => {
         clearTimeout(t);
         resolve();
@@ -113,13 +376,15 @@ export default function App() {
       try {
         wsRef.current?.close();
       } catch {
-        // ignore
+        /* ignore */
       }
     };
   }, []);
 
   const appendLog = (text) => {
-    setLog((prev) => [`${new Date().toLocaleTimeString()}  ${text}`, ...prev].slice(0, 20));
+    setLog((prev) =>
+      [`${new Date().toLocaleTimeString()}  ${text}`, ...prev].slice(0, 30)
+    );
   };
 
   const normalizeObsFromWs = (packetData) => {
@@ -150,7 +415,7 @@ export default function App() {
         drug_id_2: ids[1] || "",
         target_drug_id: ids[0] || "",
       }));
-      appendLog(`Reset task=${taskId}`);
+      appendLog(`Reset — ${TASK_LABEL_MAP[taskId] || taskId}`);
     } catch (err) {
       appendLog(`Reset failed: ${err.message}`);
     } finally {
@@ -206,7 +471,9 @@ export default function App() {
       const data = msg?.data || {};
       const normalized = normalizeObsFromWs(data);
       setObs(normalized);
-      appendLog(`Step: ${payload.action_type} -> reward=${data.reward ?? 0}`);
+      const label = ACTION_LABELS[payload.action_type] || payload.action_type;
+      const rwd = data.reward ?? 0;
+      appendLog(`${label} → reward: ${Number(rwd).toFixed(3)}`);
     } catch (err) {
       appendLog(`Step failed: ${err.message}`);
     } finally {
@@ -222,7 +489,9 @@ export default function App() {
     setLoading(true);
     try {
       const data = await apiPost("/agent/suggest", { observation: obs });
-      appendLog(`AI suggestion: ${data.action.action_type}`);
+      const label =
+        ACTION_LABELS[data.action.action_type] || data.action.action_type;
+      appendLog(`AI suggests: ${label}`);
       await handleStep(data.action);
     } catch (err) {
       appendLog(`AI suggestion failed: ${err.message}`);
@@ -231,156 +500,450 @@ export default function App() {
     }
   };
 
+  const formatDrugName = (drugId) => {
+    if (!drugId) return "";
+    return drugId
+      .replace(/^DRUG_/, "")
+      .replace(/_/g, " ")
+      .replace(/\b\w/g, (c) => c.toUpperCase());
+  };
+
+  const currentRisk = obs?.metadata?.current_risk;
+  const baselineRisk = obs?.metadata?.baseline_risk;
+
   return (
     <div className="shell">
       <div className="bg-orb orb-a" />
       <div className="bg-orb orb-b" />
 
+      {/* Spotlight Guide */}
+      {showGuide && (
+        <SpotlightGuide
+          step={guideStep}
+          steps={GUIDE_STEPS}
+          onNext={() => setGuideStep((s) => Math.min(s + 1, GUIDE_STEPS.length - 1))}
+          onPrev={() => setGuideStep((s) => Math.max(0, s - 1))}
+          onClose={() => setShowGuide(false)}
+        />
+      )}
+
       <div className="container">
-      <header className="topbar glass">
-        <div className="title-wrap">
-          <h1>Polypharmacy Control Center</h1>
-          <p>Metaverse Clinical Ops Console</p>
-        </div>
-        <div className={`status-chip ${hasValidEpisode ? "live" : "idle"}`}>
-          {hasValidEpisode ? "Session Live" : "Waiting for reset"}
-        </div>
-        <div className="actions">
-          <select value={taskId} onChange={(e) => setTaskId(e.target.value)}>
-            {TASKS.map((t) => (
-              <option key={t} value={t}>
-                {t}
-              </option>
-            ))}
-          </select>
-          <button onClick={handleReset} disabled={loading}>
-            Reset Episode
-          </button>
-          <button className="secondary" onClick={askAi} disabled={!hasValidEpisode || isDone || loading}>
-            Ask AI + Auto Step
-          </button>
-        </div>
-      </header>
-
-      <main className="layout">
-        <section className="panel glass panel-wide">
-          <h2>Episode</h2>
-          {hasValidEpisode ? (
-            <div className="kpi-grid">
-              <div><span>Episode</span><strong>{obs.episode_id}</strong></div>
-              <div><span>Task</span><strong>{obs.task_id}</strong></div>
-              <div><span>Age / Sex</span><strong>{obs.age} / {obs.sex}</strong></div>
-              <div><span>Step</span><strong>{obs.step_index}</strong></div>
-              <div><span>Query budget</span><strong>{obs.remaining_query_budget}</strong></div>
-              <div><span>Intervention budget</span><strong>{obs.remaining_intervention_budget}</strong></div>
+        <header className="topbar glass" data-guide="topbar">
+          <div className="title-wrap">
+            <h1>PolypharmacyEnv</h1>
+            <p>Elderly Medication Safety — Powered by Neural Bandits</p>
+          </div>
+          <div className="topbar-right">
+            <div className={`status-chip ${hasValidEpisode ? "live" : "idle"}`}>
+              {hasValidEpisode
+                ? isDone
+                  ? "Episode Complete"
+                  : "Session Live"
+                : "Ready"}
             </div>
-          ) : (
-            <p className="muted">Start with Reset Episode. Until then, step actions are blocked.</p>
-          )}
-          {noBudgetsLeft && (
-            <p className="muted budget-note">Query and intervention budgets are exhausted. Finish review to get final score.</p>
-          )}
-          {isDone && (
-            <p className="muted budget-note">
-              Episode complete
-              {finalScore !== null ? ` • final score: ${finalScore.toFixed(3)}` : ""}.
-              Click Reset Episode to start a new case.
-            </p>
-          )}
-        </section>
-
-        <section className="panel glass">
-          <h2>Action Console</h2>
-          <div className="action-row">
-            <label>Action type</label>
-            <select
-              value={action.action_type}
-              onChange={(e) => setAction((a) => ({ ...a, action_type: e.target.value }))}
+            <button
+              className="guide-trigger"
+              onClick={() => {
+                setGuideStep(0);
+                setShowGuide(true);
+              }}
+              title="Open guided walkthrough"
             >
-              <option value="query_ddi">query_ddi</option>
-              <option value="propose_intervention">propose_intervention</option>
-              <option value="finish_review">finish_review</option>
+              ?
+            </button>
+          </div>
+          <div className="actions" data-guide="task-selector">
+            <select value={taskId} onChange={(e) => setTaskId(e.target.value)}>
+              {TASKS.map((t) => (
+                <option key={t.id} value={t.id}>
+                  {t.label}
+                </option>
+              ))}
             </select>
+            <button onClick={handleReset} disabled={loading}>
+              Reset Episode
+            </button>
+            <button
+              className="secondary"
+              onClick={askAi}
+              disabled={!hasValidEpisode || isDone || loading}
+            >
+              Ask AI + Auto Step
+            </button>
           </div>
+        </header>
 
-          {action.action_type === "query_ddi" && (
-            <div className="stack stack-two">
-              <input
-                placeholder="drug_id_1"
-                value={action.drug_id_1}
-                onChange={(e) => setAction((a) => ({ ...a, drug_id_1: e.target.value }))}
-              />
-              <input
-                placeholder="drug_id_2"
-                value={action.drug_id_2}
-                onChange={(e) => setAction((a) => ({ ...a, drug_id_2: e.target.value }))}
-              />
-            </div>
-          )}
+        <main className="layout">
+          {/* Episode Info */}
+          <section className="panel glass panel-wide" data-guide="episode-panel">
+            <h2>Episode Overview</h2>
+            {hasValidEpisode ? (
+              <>
+                <div className="kpi-grid">
+                  <div>
+                    <span>Episode</span>
+                    <strong>{obs.episode_id}</strong>
+                  </div>
+                  <div>
+                    <span>Task</span>
+                    <strong>{TASK_LABEL_MAP[obs.task_id] || obs.task_id}</strong>
+                  </div>
+                  <div>
+                    <span>Patient</span>
+                    <strong>
+                      Age {obs.age}, {obs.sex === "M" ? "Male" : "Female"}
+                    </strong>
+                  </div>
+                  <div>
+                    <span>Step</span>
+                    <strong>{obs.step_index}</strong>
+                  </div>
+                  <div>
+                    <span>Query Budget</span>
+                    <strong>{obs.remaining_query_budget} remaining</strong>
+                  </div>
+                  <div>
+                    <span>Intervention Budget</span>
+                    <strong>
+                      {obs.remaining_intervention_budget} remaining
+                    </strong>
+                  </div>
+                </div>
 
-          {action.action_type === "propose_intervention" && (
-            <div className="stack">
+                {currentRisk !== undefined && baselineRisk !== undefined && (
+                  <div className="risk-bar-wrap">
+                    <div className="risk-labels">
+                      <span>
+                        Baseline Risk:{" "}
+                        <strong>{Number(baselineRisk).toFixed(3)}</strong>
+                      </span>
+                      <span>
+                        Current Risk:{" "}
+                        <strong
+                          className={
+                            currentRisk < baselineRisk
+                              ? "risk-down"
+                              : "risk-same"
+                          }
+                        >
+                          {Number(currentRisk).toFixed(3)}
+                        </strong>
+                      </span>
+                    </div>
+                    <div className="risk-bar">
+                      <div
+                        className="risk-fill"
+                        style={{
+                          width: `${Math.min(currentRisk * 100, 100)}%`,
+                        }}
+                      />
+                    </div>
+                  </div>
+                )}
+
+                {obs.conditions && obs.conditions.length > 0 && (
+                  <div className="conditions-row">
+                    <span className="conditions-label">Conditions:</span>
+                    {obs.conditions.map((c) => (
+                      <span key={c} className="condition-tag">
+                        {c.replace(/_/g, " ")}
+                      </span>
+                    ))}
+                  </div>
+                )}
+              </>
+            ) : (
+              <p className="muted">
+                Select a task difficulty and click <strong>Reset Episode</strong>{" "}
+                to begin a patient case.
+              </p>
+            )}
+            {noBudgetsLeft && !isDone && (
+              <div className="budget-note">
+                All budgets exhausted. Click <strong>Finish Review</strong> to
+                receive your final score.
+              </div>
+            )}
+            {isDone && (
+              <div className="budget-note done-note">
+                Episode complete
+                {finalScore !== null
+                  ? ` — Final score: ${(finalScore * 100).toFixed(1)}%`
+                  : ""}
+                . Click <strong>Reset Episode</strong> to start a new case.
+              </div>
+            )}
+          </section>
+
+          {/* Action Console */}
+          <section className="panel glass" data-guide="action-console">
+            <h2>Action Console</h2>
+            <div className="action-row">
+              <label>Action Type</label>
               <select
-                value={action.target_drug_id}
-                onChange={(e) => setAction((a) => ({ ...a, target_drug_id: e.target.value }))}
+                value={action.action_type}
+                onChange={(e) =>
+                  setAction((a) => ({ ...a, action_type: e.target.value }))
+                }
               >
-                <option value="">Select target drug</option>
-                {medIds.map((id) => (
-                  <option key={id} value={id}>
-                    {id}
+                {Object.entries(ACTION_LABELS).map(([val, label]) => (
+                  <option key={val} value={val}>
+                    {label}
                   </option>
                 ))}
               </select>
-              <select
-                value={action.intervention_type}
-                onChange={(e) => setAction((a) => ({ ...a, intervention_type: e.target.value }))}
-              >
-                <option value="stop">stop</option>
-                <option value="dose_reduce">dose_reduce</option>
-                <option value="substitute">substitute</option>
-                <option value="add_monitoring">add_monitoring</option>
-              </select>
-              <input
-                placeholder="proposed_new_drug_id (optional)"
-                value={action.proposed_new_drug_id}
-                onChange={(e) =>
-                  setAction((a) => ({ ...a, proposed_new_drug_id: e.target.value }))
-                }
-              />
-              <input
-                placeholder="rationale (optional)"
-                value={action.rationale}
-                onChange={(e) => setAction((a) => ({ ...a, rationale: e.target.value }))}
-              />
             </div>
-          )}
-          <button onClick={() => handleStep()} disabled={!isActionValid() || loading}>
-            {noBudgetsLeft ? "Finish Review" : "Submit Step"}
-          </button>
-        </section>
-
-        <section className="panel glass">
-          <h2>Current Medications</h2>
-          <div className="med-grid">
-            {(obs?.current_medications || []).map((m) => (
-              <div key={m.drug_id} className="med-card">
-                <strong>{m.drug_id}</strong>
-                <p>{m.generic_name}</p>
-                <small>{m.dose_mg} mg • {m.atc_class}</small>
+
+            {action.action_type === "query_ddi" && (
+              <div className="stack stack-two">
+                <div className="field-group">
+                  <label>Drug 1</label>
+                  <select
+                    value={action.drug_id_1}
+                    onChange={(e) =>
+                      setAction((a) => ({ ...a, drug_id_1: e.target.value }))
+                    }
+                  >
+                    <option value="">Select drug</option>
+                    {medIds.map((id) => (
+                      <option key={id} value={id}>
+                        {formatDrugName(id)}
+                      </option>
+                    ))}
+                  </select>
+                </div>
+                <div className="field-group">
+                  <label>Drug 2</label>
+                  <select
+                    value={action.drug_id_2}
+                    onChange={(e) =>
+                      setAction((a) => ({ ...a, drug_id_2: e.target.value }))
+                    }
+                  >
+                    <option value="">Select drug</option>
+                    {medIds.map((id) => (
+                      <option key={id} value={id}>
+                        {formatDrugName(id)}
+                      </option>
+                    ))}
+                  </select>
+                </div>
               </div>
-            ))}
-          </div>
-        </section>
-
-        <section className="panel glass">
-          <h2>Event Log</h2>
-          <div className="logs">
-            {log.map((line, idx) => (
-              <div key={idx}>{line}</div>
-            ))}
-          </div>
-        </section>
-      </main>
+            )}
+
+            {action.action_type === "propose_intervention" && (
+              <div className="stack">
+                <div className="field-group">
+                  <label>Target Drug</label>
+                  <select
+                    value={action.target_drug_id}
+                    onChange={(e) =>
+                      setAction((a) => ({
+                        ...a,
+                        target_drug_id: e.target.value,
+                      }))
+                    }
+                  >
+                    <option value="">Select target drug</option>
+                    {medIds.map((id) => (
+                      <option key={id} value={id}>
+                        {formatDrugName(id)}
+                      </option>
+                    ))}
+                  </select>
+                </div>
+                <div className="field-group">
+                  <label>Intervention Type</label>
+                  <select
+                    value={action.intervention_type}
+                    onChange={(e) =>
+                      setAction((a) => ({
+                        ...a,
+                        intervention_type: e.target.value,
+                      }))
+                    }
+                  >
+                    {Object.entries(INTERVENTION_LABELS).map(([val, label]) => (
+                      <option key={val} value={val}>
+                        {label}
+                      </option>
+                    ))}
+                  </select>
+                </div>
+                <div className="field-group">
+                  <label>New Drug ID (optional, for substitution)</label>
+                  <input
+                    placeholder="Leave blank for auto-selection"
+                    value={action.proposed_new_drug_id}
+                    onChange={(e) =>
+                      setAction((a) => ({
+                        ...a,
+                        proposed_new_drug_id: e.target.value,
+                      }))
+                    }
+                  />
+                </div>
+                <div className="field-group">
+                  <label>Rationale (optional)</label>
+                  <input
+                    placeholder="e.g., High bleeding risk with concurrent warfarin"
+                    value={action.rationale}
+                    onChange={(e) =>
+                      setAction((a) => ({ ...a, rationale: e.target.value }))
+                    }
+                  />
+                </div>
+              </div>
+            )}
+
+            <button
+              className="submit-btn"
+              onClick={() => handleStep()}
+              disabled={!isActionValid() || loading}
+            >
+              {noBudgetsLeft ? "Finish Review" : "Submit Step"}
+            </button>
+          </section>
+
+          {/* Current Medications */}
+          <section className="panel glass" data-guide="medications-panel">
+            <h2>
+              Current Medications
+              {obs?.current_medications?.length
+                ? ` (${obs.current_medications.length})`
+                : ""}
+            </h2>
+            <div className="med-grid">
+              {(obs?.current_medications || []).map((m) => (
+                <div
+                  key={m.drug_id}
+                  className={`med-card ${m.is_high_risk_elderly ? "high-risk" : ""}`}
+                >
+                  <div className="med-card-header">
+                    <strong>{formatDrugName(m.drug_id)}</strong>
+                    {m.is_high_risk_elderly && (
+                      <span className="risk-badge">High Risk</span>
+                    )}
+                  </div>
+                  <p className="med-generic">{m.generic_name}</p>
+                  <div className="med-details">
+                    <span>{m.dose_mg} mg</span>
+                    <span className="med-atc">{m.atc_class}</span>
+                  </div>
+                  {m.beers_flags && m.beers_flags.length > 0 && (
+                    <div className="beers-flags">
+                      {m.beers_flags.map((f, i) => (
+                        <span key={i} className="beers-tag">
+                          {f}
+                        </span>
+                      ))}
+                    </div>
+                  )}
+                </div>
+              ))}
+            </div>
+            {(!obs?.current_medications ||
+              obs.current_medications.length === 0) && (
+              <p className="muted">No medications loaded. Reset an episode to begin.</p>
+            )}
+          </section>
+
+          {/* Interaction Queries & Interventions */}
+          {hasValidEpisode && (
+            <section className="panel glass panel-wide">
+              <div className="history-grid">
+                <div>
+                  <h3>Drug Interaction Checks ({obs?.interaction_queries?.length || 0})</h3>
+                  <div className="history-list">
+                    {(obs?.interaction_queries || []).map((q, i) => (
+                      <div
+                        key={i}
+                        className={`history-item severity-${q.severity}`}
+                      >
+                        <strong>
+                          {formatDrugName(q.drug_id_1)} +{" "}
+                          {formatDrugName(q.drug_id_2)}
+                        </strong>
+                        <span className={`severity-tag ${q.severity}`}>
+                          {q.severity}
+                        </span>
+                        {q.recommendation && (
+                          <p className="history-detail">
+                            {q.recommendation.replace(/_/g, " ")}
+                          </p>
+                        )}
+                      </div>
+                    ))}
+                    {(!obs?.interaction_queries || obs.interaction_queries.length === 0) && (
+                      <p className="muted">No queries yet.</p>
+                    )}
+                  </div>
+                </div>
+                <div>
+                  <h3>Proposed Changes ({obs?.interventions?.length || 0})</h3>
+                  <div className="history-list">
+                    {(obs?.interventions || []).map((iv, i) => (
+                      <div key={i} className="history-item intervention-item">
+                        <strong>{formatDrugName(iv.target_drug_id)}</strong>
+                        <span className="intervention-tag">
+                          {INTERVENTION_LABELS[iv.action_type] || iv.action_type}
+                        </span>
+                        {iv.proposed_new_drug_id && (
+                          <p className="history-detail">
+                            Replaced with: {formatDrugName(iv.proposed_new_drug_id)}
+                          </p>
+                        )}
+                        {iv.rationale && (
+                          <p className="history-detail">{iv.rationale}</p>
+                        )}
+                      </div>
+                    ))}
+                    {(!obs?.interventions || obs.interventions.length === 0) && (
+                      <p className="muted">No interventions yet.</p>
+                    )}
+                  </div>
+                </div>
+              </div>
+            </section>
+          )}
+
+          {/* Event Log */}
+          <section className="panel glass panel-wide" data-guide="event-log">
+            <h2>Event Log</h2>
+            <div className="logs">
+              {log.length === 0 && (
+                <div className="log-empty">
+                  Events will appear here as you interact with the environment.
+                </div>
+              )}
+              {log.map((line, idx) => (
+                <div key={idx}>{line}</div>
+              ))}
+            </div>
+          </section>
+        </main>
+
+        <footer className="app-footer">
+          <p>
+            PolypharmacyEnv — Built with{" "}
+            <a
+              href="https://github.com/meta-pytorch/OpenEnv"
+              target="_blank"
+              rel="noopener noreferrer"
+            >
+              PyTorch OpenEnv
+            </a>{" "}
+            | Based on{" "}
+            <a
+              href="https://link.springer.com/chapter/10.1007/978-3-031-36938-4_5"
+              target="_blank"
+              rel="noopener noreferrer"
+            >
+              Neural Bandits for Polypharmacy
+            </a>{" "}
+            (Larouche et al.)
+          </p>
+        </footer>
       </div>
     </div>
   );

frontend/src/styles.css

@@ -1,18 +1,22 @@
 :root {
-  --bg: #070814;
-  --bg-layer: #0a1026;
-  --panel: rgba(14, 22, 44, 0.72);
-  --panel-solid: rgba(20, 28, 52, 0.92);
-  --text: #e8f1ff;
-  --muted: #9ab2db;
-  --primary: #37d4ff;
-  --primary-2: #5a8dff;
-  --accent: #9d59ff;
-  --success: #6dfbcf;
-  --border: rgba(122, 162, 255, 0.28);
-  --line: rgba(109, 143, 225, 0.18);
-  --shadow: 0 16px 45px rgba(5, 8, 23, 0.6);
-  --shadow-strong: 0 14px 32px rgba(44, 105, 255, 0.4);
+  --bg: #0d1117;
+  --bg-layer: #0d1117;
+  --panel: #161b22;
+  --panel-solid: #1c2333;
+  --text: #e6edf3;
+  --muted: #8b949e;
+  --primary: #58a6ff;
+  --primary-2: #79c0ff;
+  --accent: #58a6ff;
+  --success: #3fb950;
+  --danger: #f85149;
+  --warning: #d29922;
+  --border: rgba(48, 54, 61, 0.7);
+  --line: rgba(48, 54, 61, 0.5);
+  --shadow: 0 1px 3px rgba(0, 0, 0, 0.12), 0 4px 12px rgba(0, 0, 0, 0.08);
+  --shadow-strong: 0 4px 16px rgba(0, 0, 0, 0.2);
+  --radius: 12px;
+  --radius-sm: 8px;
 }
 
 * {
@@ -22,187 +26,253 @@
 body {
   margin: 0;
   color: var(--text);
-  font-family: "Segoe UI", "SF Pro Text", "Helvetica Neue", sans-serif;
-  background:
-    radial-gradient(circle at 8% 12%, rgba(121, 87, 255, 0.22), transparent 38%),
-    radial-gradient(circle at 88% 20%, rgba(59, 204, 255, 0.26), transparent 34%),
-    radial-gradient(circle at 50% 100%, rgba(43, 128, 255, 0.26), transparent 40%),
-    linear-gradient(145deg, var(--bg) 0%, var(--bg-layer) 60%, #04060f 100%);
-  background-attachment: fixed;
+  font-family: "Inter", -apple-system, "Segoe UI", "Helvetica Neue", sans-serif;
+  background: var(--bg);
+  line-height: 1.55;
+  -webkit-font-smoothing: antialiased;
+  -moz-osx-font-smoothing: grayscale;
 }
 
 .shell {
   min-height: 100vh;
   position: relative;
   overflow: hidden;
-  padding: 24px 16px 34px;
+  padding: 20px 24px 40px;
 }
 
 .container {
-  width: min(1320px, 100%);
+  width: min(1400px, 100%);
   margin: 0 auto;
   position: relative;
   z-index: 2;
 }
 
+/* Background orbs - subtle and muted for a professional look */
 .bg-orb {
   position: absolute;
   border-radius: 50%;
   pointer-events: none;
-  opacity: 0.9;
-  filter: blur(18px);
+  opacity: 0.3;
+  filter: blur(80px);
 }
 
 .orb-a {
-  width: min(46vw, 530px);
+  width: min(42vw, 500px);
   aspect-ratio: 1 / 1;
-  right: -9%;
-  top: -10%;
-  background: radial-gradient(circle, rgba(52, 203, 255, 0.35), rgba(52, 203, 255, 0.04) 70%);
+  right: -8%;
+  top: -8%;
+  background: radial-gradient(circle, rgba(88, 166, 255, 0.12), transparent 70%);
 }
 
 .orb-b {
-  width: min(40vw, 460px);
+  width: min(36vw, 420px);
   aspect-ratio: 1 / 1;
-  left: -9%;
-  bottom: -15%;
-  background: radial-gradient(circle, rgba(160, 102, 255, 0.3), rgba(160, 102, 255, 0.06) 72%);
+  left: -8%;
+  bottom: -12%;
+  background: radial-gradient(circle, rgba(88, 166, 255, 0.08), transparent 72%);
 }
 
+/* ── Panels ──────────────────────────────────────────────────── */
+
 .glass {
-  background:
-    linear-gradient(180deg, rgba(255, 255, 255, 0.06), rgba(255, 255, 255, 0.01)),
-    var(--panel);
+  background: var(--panel);
   border: 1px solid var(--border);
   box-shadow: var(--shadow);
-  backdrop-filter: blur(12px);
 }
 
+/* ── Top Bar ─────────────────────────────────────────────────── */
+
 .topbar {
-  border-radius: 24px;
-  padding: clamp(14px, 2vw, 20px);
-  display: grid;
-  gap: 12px 16px;
-  grid-template-columns: minmax(220px, 1.2fr) auto minmax(280px, 1fr);
+  border-radius: var(--radius);
+  padding: 14px 24px;
+  display: flex;
   align-items: center;
+  gap: 16px;
+  flex-wrap: wrap;
+}
+
+.title-wrap {
+  flex: 1;
+  min-width: 200px;
 }
 
 .title-wrap h1 {
   margin: 0;
-  font-size: clamp(1.15rem, 2.2vw, 1.95rem);
-  letter-spacing: 0.02em;
-  text-transform: uppercase;
-  text-shadow: 0 0 16px rgba(106, 192, 255, 0.3);
+  font-size: clamp(1rem, 1.8vw, 1.35rem);
+  font-weight: 700;
+  letter-spacing: -0.01em;
+  color: var(--text);
+  background: none;
+  -webkit-background-clip: unset;
+  -webkit-text-fill-color: unset;
+  background-clip: unset;
 }
 
 .title-wrap p {
-  margin: 6px 0 0;
-  font-size: 0.84rem;
+  margin: 2px 0 0;
+  font-size: 0.8rem;
   color: var(--muted);
-  letter-spacing: 0.03em;
-  text-transform: uppercase;
+  letter-spacing: 0.01em;
+}
+
+.topbar-right {
+  display: flex;
+  align-items: center;
+  gap: 10px;
 }
 
 .status-chip {
-  justify-self: center;
-  padding: 7px 14px;
+  padding: 5px 14px;
   border-radius: 999px;
   font-size: 0.72rem;
-  font-weight: 700;
-  letter-spacing: 0.08em;
+  font-weight: 600;
+  letter-spacing: 0.04em;
   text-transform: uppercase;
   border: 1px solid transparent;
+  white-space: nowrap;
 }
 
 .status-chip.live {
-  color: #052c24;
-  background: linear-gradient(90deg, rgba(126, 255, 220, 0.9), rgba(84, 244, 196, 0.95));
-  box-shadow: 0 0 14px rgba(96, 244, 198, 0.36);
+  color: #ffffff;
+  background: var(--success);
+  box-shadow: none;
 }
 
 .status-chip.idle {
-  color: #d8e8ff;
-  border-color: rgba(117, 186, 255, 0.48);
-  background: rgba(60, 106, 198, 0.25);
+  color: var(--muted);
+  border-color: var(--border);
+  background: rgba(48, 54, 61, 0.3);
+}
+
+.guide-trigger {
+  width: 34px !important;
+  height: 34px;
+  min-height: 34px;
+  padding: 0;
+  border-radius: 50% !important;
+  font-size: 0.95rem;
+  font-weight: 700;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  background: rgba(88, 166, 255, 0.1);
+  border: 1px solid rgba(88, 166, 255, 0.25);
+  color: var(--primary);
+  box-shadow: none;
+}
+
+.guide-trigger:hover {
+  background: rgba(88, 166, 255, 0.2);
+  box-shadow: none;
 }
 
 .actions {
   display: flex;
-  justify-content: flex-end;
+  align-items: center;
   flex-wrap: wrap;
   gap: 10px;
 }
 
+/* ── Form Controls ───────────────────────────────────────────── */
+
 button,
 select,
 input {
   width: 100%;
   min-height: 42px;
-  border-radius: 12px;
+  border-radius: var(--radius-sm);
   border: 1px solid var(--border);
-  font-size: 0.92rem;
-  padding: 10px 12px;
+  font-size: 0.9rem;
+  padding: 10px 14px;
   color: var(--text);
-  background: rgba(11, 19, 38, 0.84);
+  background: #0d1117;
+  font-family: inherit;
+  transition: border-color 150ms ease, box-shadow 150ms ease, background 150ms ease;
 }
 
-select,
-input {
-  transition: border-color 120ms ease, box-shadow 120ms ease;
+select:hover,
+input:hover {
+  border-color: rgba(139, 148, 158, 0.5);
 }
 
 select:focus,
 input:focus {
   outline: none;
-  border-color: rgba(119, 200, 255, 0.88);
-  box-shadow: 0 0 0 2px rgba(95, 187, 255, 0.18);
+  border-color: var(--primary);
+  box-shadow: 0 0 0 3px rgba(88, 166, 255, 0.15);
+  background: #0d1117;
+}
+
+select {
+  cursor: pointer;
+  appearance: auto;
 }
 
 button {
   cursor: pointer;
   border: 0;
   width: auto;
-  font-weight: 700;
-  letter-spacing: 0.02em;
-  background: linear-gradient(135deg, var(--primary), var(--primary-2) 55%, var(--accent));
-  box-shadow: var(--shadow-strong);
-  transition: transform 140ms ease, filter 140ms ease, box-shadow 140ms ease;
+  font-weight: 600;
+  letter-spacing: 0.01em;
+  white-space: nowrap;
+  background: var(--primary);
+  color: #ffffff;
+  box-shadow: none;
+  transition: background 150ms ease, transform 100ms ease, box-shadow 150ms ease;
 }
 
 button:hover {
+  background: #79c0ff;
   transform: translateY(-1px);
-  filter: brightness(1.04);
-  box-shadow: 0 18px 32px rgba(50, 141, 255, 0.48);
+  filter: none;
+  box-shadow: 0 2px 8px rgba(88, 166, 255, 0.25);
 }
 
 button:active {
   transform: translateY(0);
+  background: #4090e0;
 }
 
 button.secondary {
-  background: linear-gradient(135deg, rgba(95, 185, 255, 0.9), rgba(154, 102, 255, 0.86));
+  background: rgba(88, 166, 255, 0.15);
+  color: var(--primary);
+  border: 1px solid rgba(88, 166, 255, 0.3);
+}
+
+button.secondary:hover {
+  background: rgba(88, 166, 255, 0.25);
 }
 
 button:disabled {
-  opacity: 0.56;
+  opacity: 0.4;
   cursor: not-allowed;
-  filter: grayscale(0.2);
+  filter: none;
   box-shadow: none;
   transform: none;
 }
 
+.submit-btn {
+  width: 100%;
+  margin-top: 10px;
+  min-height: 44px;
+  font-size: 0.92rem;
+}
+
+/* ── Layout ──────────────────────────────────────────────────── */
+
 .layout {
-  margin-top: 16px;
+  margin-top: 20px;
   display: grid;
-  gap: 14px;
-  grid-template-columns: 1.12fr 0.88fr;
+  gap: 20px;
+  grid-template-columns: 1.1fr 0.9fr;
   align-items: start;
 }
 
 .panel {
-  border-radius: 20px;
-  padding: clamp(14px, 1.8vw, 20px);
+  border-radius: var(--radius);
+  padding: 24px;
   position: relative;
+  overflow: hidden;
 }
 
 .panel::after {
@@ -219,151 +289,576 @@ button:disabled {
 }
 
 .panel h2 {
-  margin: 0 0 12px;
-  font-size: 1rem;
-  font-weight: 700;
+  margin: 0 0 16px;
+  font-size: 0.82rem;
+  font-weight: 600;
   letter-spacing: 0.05em;
   text-transform: uppercase;
+  color: var(--muted);
+}
+
+.panel h3 {
+  margin: 0 0 12px;
+  font-size: 0.8rem;
+  font-weight: 600;
+  letter-spacing: 0.04em;
+  text-transform: uppercase;
+  color: var(--muted);
 }
 
+/* ── KPI Grid ────────────────────────────────────────────────── */
+
 .kpi-grid {
   display: grid;
-  gap: 10px;
-  grid-template-columns: repeat(3, minmax(0, 1fr));
+  gap: 12px;
+  grid-template-columns: repeat(3, 1fr);
 }
 
 .kpi-grid div {
-  border-radius: 13px;
+  border-radius: var(--radius-sm);
   border: 1px solid var(--border);
   background: var(--panel-solid);
-  padding: 11px 12px;
+  padding: 16px 18px;
+  overflow: hidden;
 }
 
 .kpi-grid span {
   display: block;
-  margin-bottom: 4px;
+  margin-bottom: 6px;
   font-size: 0.72rem;
   color: var(--muted);
   text-transform: uppercase;
-  letter-spacing: 0.05em;
+  letter-spacing: 0.06em;
 }
 
 .kpi-grid strong {
-  font-size: 1.06rem;
-  line-height: 1.2;
+  font-size: 1.05rem;
+  font-weight: 600;
+  line-height: 1.3;
+  word-break: break-word;
+  overflow-wrap: break-word;
+  color: var(--text);
+}
+
+/* ── Risk Bar ────────────────────────────────────────────────── */
+
+.risk-bar-wrap {
+  margin-top: 16px;
+}
+
+.risk-labels {
+  display: flex;
+  justify-content: space-between;
+  font-size: 0.8rem;
+  color: var(--muted);
+  margin-bottom: 8px;
 }
 
+.risk-labels strong {
+  font-size: 0.85rem;
+  color: var(--text);
+}
+
+.risk-down {
+  color: var(--success) !important;
+}
+
+.risk-same {
+  color: var(--warning) !important;
+}
+
+.risk-bar {
+  height: 6px;
+  background: rgba(48, 54, 61, 0.5);
+  border-radius: 3px;
+  overflow: hidden;
+}
+
+.risk-fill {
+  height: 100%;
+  background: linear-gradient(90deg, var(--success), var(--warning), var(--danger));
+  border-radius: 3px;
+  transition: width 300ms ease;
+}
+
+/* ── Conditions ──────────────────────────────────────────────── */
+
+.conditions-row {
+  margin-top: 14px;
+  display: flex;
+  flex-wrap: wrap;
+  align-items: center;
+  gap: 8px;
+}
+
+.conditions-label {
+  font-size: 0.78rem;
+  color: var(--muted);
+  text-transform: uppercase;
+  letter-spacing: 0.04em;
+  margin-right: 4px;
+}
+
+.condition-tag {
+  font-size: 0.75rem;
+  padding: 4px 12px;
+  border-radius: 999px;
+  background: rgba(88, 166, 255, 0.1);
+  border: 1px solid rgba(88, 166, 255, 0.2);
+  color: var(--primary);
+  text-transform: capitalize;
+  white-space: nowrap;
+}
+
+/* ── Action Console ──────────────────────────────────────────── */
+
 .action-row,
 .stack {
   display: grid;
-  gap: 10px;
-  margin-bottom: 12px;
+  gap: 12px;
+  margin-bottom: 14px;
 }
 
-.action-row label {
+.action-row label,
+.field-group label {
+  display: block;
   color: var(--muted);
   font-size: 0.78rem;
-  letter-spacing: 0.05em;
+  letter-spacing: 0.04em;
   text-transform: uppercase;
+  margin-bottom: 6px;
+}
+
+.field-group {
+  display: flex;
+  flex-direction: column;
 }
 
 .stack-two {
-  grid-template-columns: repeat(2, minmax(0, 1fr));
+  grid-template-columns: 1fr 1fr;
 }
 
+/* ── Medication Cards ────────────────────────────────────────── */
+
 .med-grid {
   display: grid;
-  grid-template-columns: repeat(3, minmax(0, 1fr));
-  gap: 10px;
-  max-height: 430px;
-  overflow: auto;
+  grid-template-columns: repeat(auto-fill, minmax(210px, 1fr));
+  gap: 12px;
+  max-height: 480px;
+  overflow-y: auto;
   padding-right: 4px;
 }
 
 .med-card {
-  border-radius: 14px;
+  border-radius: var(--radius-sm);
   border: 1px solid var(--border);
   background: var(--panel-solid);
-  padding: 11px 12px;
-  transition: transform 130ms ease, border-color 130ms ease;
+  padding: 16px 18px;
+  transition: border-color 150ms ease, background 150ms ease;
+  overflow: hidden;
 }
 
 .med-card:hover {
-  transform: translateY(-1px);
-  border-color: rgba(109, 224, 255, 0.72);
+  border-color: rgba(88, 166, 255, 0.4);
+  background: #1f2937;
+}
+
+.med-card.high-risk {
+  border-color: rgba(248, 81, 73, 0.35);
+}
+
+.med-card-header {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 8px;
+}
+
+.med-card-header strong {
+  font-size: 0.92rem;
+  font-weight: 600;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+  color: var(--text);
+}
+
+.risk-badge {
+  font-size: 0.65rem;
+  padding: 3px 9px;
+  border-radius: 999px;
+  background: rgba(248, 81, 73, 0.12);
+  border: 1px solid rgba(248, 81, 73, 0.3);
+  color: var(--danger);
+  text-transform: uppercase;
+  letter-spacing: 0.04em;
+  font-weight: 600;
+  white-space: nowrap;
+  flex-shrink: 0;
 }
 
-.med-card p {
-  margin: 6px 0 4px;
+.med-generic {
+  margin: 6px 0;
   color: var(--muted);
+  font-size: 0.84rem;
   text-transform: capitalize;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
 }
 
-.med-card small {
-  color: #c7d9ff;
+.med-details {
+  display: flex;
+  gap: 10px;
+  font-size: 0.8rem;
+  color: #8b949e;
 }
 
-.logs {
+.med-atc {
+  color: var(--primary);
+  font-weight: 600;
+}
+
+.beers-flags {
+  margin-top: 8px;
+  display: flex;
+  flex-wrap: wrap;
+  gap: 5px;
+}
+
+.beers-tag {
+  font-size: 0.68rem;
+  padding: 3px 9px;
+  border-radius: 999px;
+  background: rgba(210, 153, 34, 0.1);
+  border: 1px solid rgba(210, 153, 34, 0.25);
+  color: var(--warning);
+  text-transform: capitalize;
+}
+
+/* ── History Grid ────────────────────────────────────────────── */
+
+.history-grid {
+  display: grid;
+  grid-template-columns: 1fr 1fr;
+  gap: 24px;
+}
+
+.history-list {
+  display: flex;
+  flex-direction: column;
+  gap: 10px;
   max-height: 300px;
-  overflow: auto;
+  overflow-y: auto;
+}
+
+.history-item {
+  border-radius: var(--radius-sm);
+  border: 1px solid var(--border);
+  background: var(--panel-solid);
+  padding: 14px 16px;
+}
+
+.history-item strong {
+  font-size: 0.88rem;
+  font-weight: 600;
+  display: block;
+  margin-bottom: 6px;
+  color: var(--text);
+}
+
+.history-detail {
+  margin: 6px 0 0;
+  font-size: 0.82rem;
+  color: var(--muted);
+  text-transform: capitalize;
+}
+
+.severity-tag, .intervention-tag {
+  display: inline-block;
+  font-size: 0.7rem;
+  padding: 3px 10px;
+  border-radius: 999px;
+  font-weight: 600;
+  text-transform: uppercase;
+  letter-spacing: 0.04em;
+}
+
+.severity-tag.severe {
+  background: rgba(248, 81, 73, 0.12);
+  border: 1px solid rgba(248, 81, 73, 0.3);
+  color: var(--danger);
+}
+
+.severity-tag.moderate {
+  background: rgba(210, 153, 34, 0.12);
+  border: 1px solid rgba(210, 153, 34, 0.3);
+  color: var(--warning);
+}
+
+.severity-tag.mild {
+  background: rgba(63, 185, 80, 0.12);
+  border: 1px solid rgba(63, 185, 80, 0.3);
+  color: var(--success);
+}
+
+.severity-tag.none {
+  background: rgba(48, 54, 61, 0.3);
+  border: 1px solid var(--border);
+  color: var(--muted);
+}
+
+.intervention-tag {
+  background: rgba(88, 166, 255, 0.1);
+  border: 1px solid rgba(88, 166, 255, 0.25);
+  color: var(--primary);
+}
+
+.severity-severe .history-item {
+  border-color: rgba(248, 81, 73, 0.2);
+}
+
+.severity-moderate .history-item {
+  border-color: rgba(210, 153, 34, 0.2);
+}
+
+/* ── Event Log ───────────────────────────────────────────────── */
+
+.logs {
+  max-height: 280px;
+  overflow-y: auto;
   padding-right: 4px;
-  display: grid;
-  gap: 7px;
+  display: flex;
+  flex-direction: column;
+  gap: 8px;
   font-size: 0.84rem;
-  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, monospace;
+  font-family: "JetBrains Mono", ui-monospace, SFMono-Regular, Menlo, Monaco, monospace;
 }
 
 .logs div {
-  border-radius: 10px;
+  border-radius: var(--radius-sm);
   border: 1px solid var(--border);
-  background: rgba(10, 16, 31, 0.84);
-  padding: 8px 10px;
-  color: #dbebff;
+  background: #0d1117;
+  padding: 12px 16px;
+  color: var(--text);
+  overflow: hidden;
+  text-overflow: ellipsis;
+  word-break: break-word;
+  line-height: 1.5;
+}
+
+.log-empty {
+  color: var(--muted);
+  font-family: inherit;
+  font-style: italic;
 }
 
+/* ── Helper Text ─────────────────────────────────────────────── */
+
 .muted {
   margin: 0;
   color: var(--muted);
+  font-size: 0.88rem;
 }
 
 .budget-note {
-  margin-top: 10px;
+  margin-top: 14px;
   border: 1px solid var(--border);
-  border-radius: 12px;
-  padding: 10px 12px;
-  background: rgba(13, 22, 42, 0.82);
+  border-radius: var(--radius-sm);
+  padding: 14px 18px;
+  background: var(--panel-solid);
+  font-size: 0.88rem;
+  color: var(--muted);
+}
+
+.done-note {
+  border-color: rgba(63, 185, 80, 0.3);
+  color: var(--success);
+}
+
+/* ── Footer ──────────────────────────────────────────────────── */
+
+.app-footer {
+  margin-top: 28px;
+  text-align: center;
+}
+
+.app-footer p {
+  font-size: 0.8rem;
+  color: var(--muted);
+  opacity: 0.6;
+}
+
+.app-footer a {
+  color: var(--primary);
+  text-decoration: none;
+}
+
+.app-footer a:hover {
+  text-decoration: underline;
 }
 
+/* ── Spotlight Guide ─────────────────────────────────────────── */
+
+.spotlight-overlay {
+  position: fixed;
+  inset: 0;
+  z-index: 100;
+  pointer-events: none;
+}
+
+.spotlight-svg {
+  position: fixed;
+  inset: 0;
+  z-index: 100;
+  pointer-events: auto;
+}
+
+.spotlight-ring {
+  position: fixed;
+  z-index: 101;
+  border: 2px solid var(--primary);
+  border-radius: var(--radius);
+  box-shadow: 0 0 0 4px rgba(88, 166, 255, 0.1);
+  pointer-events: none;
+  transition: top 350ms ease, left 350ms ease, width 350ms ease, height 350ms ease;
+}
+
+.spotlight-tooltip {
+  position: fixed;
+  z-index: 102;
+  border-radius: var(--radius);
+  padding: 24px;
+  pointer-events: auto;
+  animation: tooltipIn 250ms ease;
+}
+
+@keyframes tooltipIn {
+  from {
+    opacity: 0;
+    transform: translateY(8px);
+  }
+  to {
+    opacity: 1;
+    transform: translateY(0);
+  }
+}
+
+.spotlight-tooltip-header {
+  display: flex;
+  align-items: flex-start;
+  justify-content: space-between;
+  gap: 12px;
+  margin-bottom: 12px;
+}
+
+.spotlight-tooltip-header h3 {
+  margin: 0;
+  font-size: 1.05rem;
+  font-weight: 600;
+  color: var(--text);
+  text-transform: none;
+  letter-spacing: 0;
+  line-height: 1.3;
+}
+
+.guide-counter {
+  font-size: 0.72rem;
+  color: var(--muted);
+  padding: 4px 10px;
+  border-radius: 999px;
+  background: rgba(48, 54, 61, 0.4);
+  border: 1px solid var(--border);
+  white-space: nowrap;
+  flex-shrink: 0;
+}
+
+.spotlight-tooltip-body {
+  margin-bottom: 18px;
+  font-size: 0.88rem;
+  line-height: 1.65;
+  color: var(--muted);
+}
+
+.spotlight-tooltip-body p {
+  margin: 0 0 5px;
+}
+
+.spotlight-tooltip-body p:empty {
+  height: 4px;
+}
+
+.spotlight-tooltip-footer {
+  display: flex;
+  gap: 8px;
+  justify-content: flex-end;
+}
+
+.guide-btn {
+  padding: 8px 18px !important;
+  font-size: 0.84rem !important;
+  border-radius: var(--radius-sm) !important;
+}
+
+.guide-dots {
+  display: flex;
+  justify-content: center;
+  gap: 6px;
+  margin-top: 14px;
+}
+
+.dot {
+  width: 7px;
+  height: 7px;
+  border-radius: 50%;
+  background: rgba(48, 54, 61, 0.6);
+  transition: background 150ms ease;
+}
+
+.dot.active {
+  background: var(--primary);
+  box-shadow: none;
+}
+
+/* ── Responsive ──────────────────────────────────────────────── */
+
 @media (max-width: 1180px) {
   .layout {
     grid-template-columns: 1fr;
   }
 
   .topbar {
-    grid-template-columns: 1fr;
+    flex-direction: column;
+    align-items: flex-start;
   }
 
-  .status-chip {
-    justify-self: start;
+  .topbar-right {
+    align-self: flex-start;
   }
 
   .actions {
+    width: 100%;
     justify-content: flex-start;
   }
+
+  .history-grid {
+    grid-template-columns: 1fr;
+  }
 }
 
 @media (max-width: 760px) {
   .shell {
-    padding: 14px 10px 24px;
+    padding: 12px 12px 24px;
   }
 
   .topbar,
   .panel {
-    border-radius: 16px;
+    border-radius: var(--radius-sm);
+    padding: 16px 18px;
   }
 
   .actions {
-    width: 100%;
+    flex-direction: column;
   }
 
   .actions button,
@@ -371,13 +866,53 @@ button:disabled {
     width: 100%;
   }
 
-  .kpi-grid,
-  .med-grid,
+  .kpi-grid {
+    grid-template-columns: 1fr 1fr;
+  }
+
+  .med-grid {
+    grid-template-columns: 1fr;
+  }
+
   .stack-two {
     grid-template-columns: 1fr;
   }
 
+  .guide-modal {
+    padding: 20px;
+  }
+
+  .spotlight-tooltip {
+    left: 10px !important;
+    right: 10px !important;
+    max-width: calc(100vw - 20px) !important;
+  }
+
+  .guide-footer,
+  .spotlight-tooltip-footer {
+    flex-direction: column;
+  }
+
   .logs {
-    max-height: 240px;
+    max-height: 200px;
   }
 }
+
+/* ── Scrollbar ───────────────────────────────────────────────── */
+
+::-webkit-scrollbar {
+  width: 6px;
+}
+
+::-webkit-scrollbar-track {
+  background: transparent;
+}
+
+::-webkit-scrollbar-thumb {
+  background: rgba(48, 54, 61, 0.6);
+  border-radius: 3px;
+}
+
+::-webkit-scrollbar-thumb:hover {
+  background: rgba(139, 148, 158, 0.4);
+}

train_bandit.py

@@ -0,0 +1,381 @@
+#!/usr/bin/env python3
+"""OptimNeuralTS training -- Neural Bandit search for dangerous polypharmacies.
+
+Implements the training pipeline from:
+  Larouche et al., "Neural Bandits for Data Mining: Searching for Dangerous Polypharmacy"
+  https://link.springer.com/chapter/10.1007/978-3-031-36938-4_5
+
+This script:
+  1. Generates a synthetic dataset of drug combinations with simulated Relative Risk (RR)
+  2. Runs OptimNeuralTS: warm-up -> NeuralTS+DE exploration -> ensemble building
+  3. Evaluates the ensemble's ability to detect Potentially Inappropriate Polypharmacies (PIPs)
+  4. Saves the trained ensemble model
+
+Usage:
+    python train_bandit.py --total-steps 1000 --warmup-steps 200
+    python train_bandit.py --total-steps 3000 --warmup-steps 500 --eval-every 100
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import sys
+import time
+from itertools import combinations
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import torch
+
+_BACKEND_SRC = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "backend", "src"
+)
+sys.path.insert(0, _BACKEND_SRC)
+
+from polypharmacy_env.neural_bandits import NeuralTS, OptimNeuralTS, nearest_neighbor_hamming  # noqa: E402
+from polypharmacy_env.data_loader import load_drug_metadata, load_ddi_rules  # noqa: E402
+
+
+# ---------------------------------------------------------------------------
+# Synthetic RR data generation (follows paper Section 4.1)
+# ---------------------------------------------------------------------------
+
+def generate_synthetic_dataset(
+    n_drugs: int = 33,
+    n_combinations: int = 5000,
+    n_dangerous_patterns: int = 10,
+    rr_threshold: float = 1.1,
+    noise_std: float = 0.1,
+    seed: int = 42,
+) -> Dict[str, Any]:
+    """Generate synthetic drug combination data with ground-truth RRs.
+
+    Follows the paper's data generation process:
+    - Generate dangerous patterns (binomial)
+    - For each combination, compute similarity to nearest pattern
+    - Assign RR proportional to similarity (if overlapping) or from N(mu, sigma) if disjoint
+    """
+    rng = random.Random(seed)
+    torch.manual_seed(seed)
+
+    # Generate dangerous patterns (multi-hot vectors)
+    patterns = []
+    for _ in range(n_dangerous_patterns):
+        # Each drug has ~30% chance of being in the pattern (smaller patterns)
+        p = torch.zeros(n_drugs)
+        n_active = rng.randint(2, max(3, n_drugs // 8))
+        indices = rng.sample(range(n_drugs), n_active)
+        for idx in indices:
+            p[idx] = 1.0
+        patterns.append(p)
+
+    # Generate distinct drug combinations
+    combos = []
+    combo_set = set()
+    while len(combos) < n_combinations:
+        n_active = rng.randint(2, min(8, n_drugs))
+        indices = tuple(sorted(rng.sample(range(n_drugs), n_active)))
+        if indices not in combo_set:
+            combo_set.add(indices)
+            vec = torch.zeros(n_drugs)
+            for idx in indices:
+                vec[idx] = 1.0
+            combos.append(vec)
+
+    # Compute RR for each combination based on Hamming distance to nearest pattern
+    rrs = []
+    nearest_pattern_idx = []
+    for combo in combos:
+        # Find nearest pattern (Hamming distance)
+        min_dist = float("inf")
+        best_p_idx = 0
+        for p_idx, pattern in enumerate(patterns):
+            dist = (combo != pattern).float().sum().item()
+            if dist < min_dist:
+                min_dist = dist
+                best_p_idx = p_idx
+
+        nearest_pattern_idx.append(best_p_idx)
+        pattern = patterns[best_p_idx]
+
+        # Check intersection (shared active drugs)
+        intersection = (combo * pattern).sum().item()
+        if intersection > 0:
+            # RR proportional to similarity
+            similarity = intersection / max(pattern.sum().item(), 1)
+            # Higher similarity -> higher RR
+            base_rr = 0.5 + 2.5 * similarity  # range ~[0.5, 3.0]
+            noise = rng.gauss(0, 0.15)
+            rr = max(0.1, base_rr + noise)
+        else:
+            # Disjoint: sample from neutral distribution
+            rr = max(0.1, rng.gauss(0.85, 0.2))
+
+        rrs.append(rr)
+
+    # Compute pattern RRs (patterns themselves have high RR)
+    pattern_rrs = [2.0 + rng.gauss(0, 0.3) for _ in patterns]
+
+    n_pips = sum(1 for rr in rrs if rr > rr_threshold)
+    print(f"  Generated {n_combinations} combos, {n_pips} PIPs (RR > {rr_threshold})")
+    print(f"  RR range: [{min(rrs):.3f}, {max(rrs):.3f}], mean: {sum(rrs)/len(rrs):.3f}")
+
+    return {
+        "combos": combos,
+        "rrs": rrs,
+        "patterns": patterns,
+        "pattern_rrs": pattern_rrs,
+        "n_drugs": n_drugs,
+        "n_pips": n_pips,
+        "rr_threshold": rr_threshold,
+        "noise_std": noise_std,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Training loop
+# ---------------------------------------------------------------------------
+
+def train_bandit(args: argparse.Namespace) -> None:
+    print("=" * 72)
+    print("OptimNeuralTS Training -- Neural Bandits for Polypharmacy")
+    print("=" * 72)
+
+    # Generate synthetic data
+    print("\nGenerating synthetic dataset...")
+    dataset = generate_synthetic_dataset(
+        n_drugs=args.n_drugs,
+        n_combinations=args.n_combinations,
+        n_dangerous_patterns=args.n_patterns,
+        seed=args.seed,
+    )
+    combos = dataset["combos"]
+    rrs = dataset["rrs"]
+    patterns = dataset["patterns"]
+    pattern_rrs = dataset["pattern_rrs"]
+    noise_std = dataset["noise_std"]
+    rr_threshold = dataset["rr_threshold"]
+
+    # Initialize OptimNeuralTS
+    bandit = OptimNeuralTS(
+        input_dim=args.n_drugs,
+        hidden=args.hidden_dim,
+        reg_lambda=args.reg_lambda,
+        exploration_factor=args.exploration_factor,
+        lr=args.lr,
+        train_epochs=args.train_epochs,
+        warmup_steps=args.warmup_steps,
+        total_steps=args.total_steps,
+        retrain_every=args.retrain_every,
+        de_population=args.de_population,
+        de_crossover=args.de_crossover,
+        de_weight=args.de_weight,
+        de_steps=args.de_steps,
+    )
+
+    print(f"\n  n_drugs           : {args.n_drugs}")
+    print(f"  n_combinations    : {args.n_combinations}")
+    print(f"  total_steps (T)   : {args.total_steps}")
+    print(f"  warmup_steps (τ)  : {args.warmup_steps}")
+    print(f"  DE population (N) : {args.de_population}")
+    print(f"  DE steps (S)      : {args.de_steps}")
+    print(f"  retrain_every     : {args.retrain_every}")
+    print(f"  hidden_dim        : {args.hidden_dim}")
+    print(f"  lr                : {args.lr}")
+    print("=" * 72)
+
+    t_start = time.time()
+
+    # Metrics tracking
+    step_rewards = []
+    pips_found = []
+    eval_precisions = []
+    eval_recalls = []
+    training_dataset_indices = set()
+
+    for t in range(1, args.total_steps + 1):
+        # Select action
+        idx, info = bandit.select_action(combos)
+        training_dataset_indices.add(idx)
+
+        # Observe noisy reward (RR + noise)
+        true_rr = rrs[idx]
+        noisy_rr = true_rr + random.gauss(0, noise_std)
+        reward = noisy_rr
+
+        step_rewards.append(reward)
+
+        # Update bandit
+        loss = bandit.observe(combos[idx], reward)
+
+        # Periodic evaluation
+        if t % args.eval_every == 0 or t == args.total_steps:
+            # Evaluate ensemble on ALL combinations
+            true_positives = 0
+            false_positives = 0
+            true_negatives = 0
+            false_negatives = 0
+
+            for i, combo in enumerate(combos):
+                pred = bandit.predict_risk(combo)
+                actual_pip = rrs[i] > rr_threshold
+                predicted_pip = pred["is_potentially_harmful"]
+
+                if predicted_pip and actual_pip:
+                    true_positives += 1
+                elif predicted_pip and not actual_pip:
+                    false_positives += 1
+                elif not predicted_pip and actual_pip:
+                    false_negatives += 1
+                else:
+                    true_negatives += 1
+
+            precision = true_positives / max(true_positives + false_positives, 1)
+            recall = true_positives / max(true_positives + false_negatives, 1)
+            eval_precisions.append(precision)
+            eval_recalls.append(recall)
+
+            # Check dangerous pattern detection
+            patterns_found = 0
+            for p_idx, pattern in enumerate(patterns):
+                pred = bandit.predict_risk(pattern)
+                if pred["is_potentially_harmful"]:
+                    patterns_found += 1
+            pattern_ratio = patterns_found / len(patterns)
+
+            # PIPs found outside training data
+            pips_outside_train = 0
+            total_detected_pips = 0
+            for i, combo in enumerate(combos):
+                pred = bandit.predict_risk(combo)
+                if pred["is_potentially_harmful"]:
+                    total_detected_pips += 1
+                    if i not in training_dataset_indices:
+                        pips_outside_train += 1
+
+            pips_found.append(total_detected_pips)
+
+            elapsed = time.time() - t_start
+            phase = info.get("phase", "?")
+            n_ens = len(bandit.agent.ensemble_weights)
+            print(
+                f"[step {t:>5d}/{args.total_steps}] "
+                f"phase={phase}  "
+                f"precision={precision:.3f}  "
+                f"recall={recall:.3f}  "
+                f"patterns={pattern_ratio:.2f}  "
+                f"PIPs_detected={total_detected_pips}  "
+                f"outside_train={pips_outside_train}  "
+                f"ensemble={n_ens}  "
+                f"elapsed={elapsed:.1f}s"
+            )
+
+    # Save metrics
+    metrics = {
+        "algorithm": "OptimNeuralTS",
+        "n_drugs": args.n_drugs,
+        "n_combinations": args.n_combinations,
+        "total_steps": args.total_steps,
+        "warmup_steps": args.warmup_steps,
+        "n_ensemble_models": len(bandit.agent.ensemble_weights),
+        "final_precision": eval_precisions[-1] if eval_precisions else 0,
+        "final_recall": eval_recalls[-1] if eval_recalls else 0,
+        "eval_precisions": eval_precisions,
+        "eval_recalls": eval_recalls,
+        "pips_detected": pips_found,
+        "step_rewards": step_rewards,
+        "total_time_s": time.time() - t_start,
+        "hyperparameters": {
+            "hidden_dim": args.hidden_dim,
+            "lr": args.lr,
+            "reg_lambda": args.reg_lambda,
+            "exploration_factor": args.exploration_factor,
+            "de_population": args.de_population,
+            "de_crossover": args.de_crossover,
+            "de_weight": args.de_weight,
+            "de_steps": args.de_steps,
+            "train_epochs": args.train_epochs,
+            "retrain_every": args.retrain_every,
+        },
+    }
+
+    metrics_path = Path(args.metrics_file)
+    metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(metrics_path, "w") as f:
+        json.dump(metrics, f, indent=2)
+    print(f"\nMetrics saved to {metrics_path}")
+
+    # Save model ensemble
+    ckpt_dir = Path(args.checkpoint_dir)
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    ckpt_path = ckpt_dir / "bandit_ensemble.pt"
+    torch.save({
+        "ensemble_weights": bandit.agent.ensemble_weights,
+        "network_state_dict": bandit.agent.network.state_dict(),
+        "U_diag": bandit.agent.U_diag,
+        "input_dim": args.n_drugs,
+        "hidden_dim": args.hidden_dim,
+        "n_steps": args.total_steps,
+    }, ckpt_path)
+    print(f"Ensemble model saved to {ckpt_path}")
+
+    print(f"\n{'='*72}")
+    print("Training complete!")
+    print(f"  Ensemble size: {len(bandit.agent.ensemble_weights)} models")
+    if eval_precisions:
+        print(f"  Final precision: {eval_precisions[-1]:.4f}")
+        print(f"  Final recall: {eval_recalls[-1]:.4f}")
+    print(f"  Total time: {time.time() - t_start:.1f}s")
+    print(f"{'='*72}")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="OptimNeuralTS training for polypharmacy PIP detection",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    # Dataset
+    p.add_argument("--n-drugs", type=int, default=33, help="Number of possible drugs")
+    p.add_argument("--n-combinations", type=int, default=5000, help="Number of distinct drug combinations")
+    p.add_argument("--n-patterns", type=int, default=10, help="Number of dangerous patterns")
+    p.add_argument("--seed", type=int, default=42, help="Random seed")
+
+    # OptimNeuralTS
+    p.add_argument("--total-steps", type=int, default=1000, help="Total bandit steps T")
+    p.add_argument("--warmup-steps", type=int, default=200, help="Warmup steps τ")
+    p.add_argument("--retrain-every", type=int, default=10, help="Retrain network every N steps")
+    p.add_argument("--hidden-dim", type=int, default=64, help="Network hidden layer size")
+    p.add_argument("--lr", type=float, default=0.01, help="Learning rate")
+    p.add_argument("--reg-lambda", type=float, default=1.0, help="Regularization λ")
+    p.add_argument("--exploration-factor", type=float, default=1.0, help="Exploration ν")
+    p.add_argument("--train-epochs", type=int, default=50, help="Epochs per retrain")
+
+    # DE
+    p.add_argument("--de-population", type=int, default=16, help="DE population size N")
+    p.add_argument("--de-crossover", type=float, default=0.9, help="DE crossover rate C")
+    p.add_argument("--de-weight", type=float, default=1.0, help="DE differential weight F")
+    p.add_argument("--de-steps", type=int, default=8, help="DE optimization steps S")
+
+    # Output
+    p.add_argument("--eval-every", type=int, default=100, help="Evaluate every N steps")
+    p.add_argument("--metrics-file", type=str, default="bandit_metrics.json", help="Metrics output path")
+    p.add_argument(
+        "--checkpoint-dir", type=str,
+        default=os.path.join(_BACKEND_SRC, "polypharmacy_env", "checkpoints"),
+        help="Model checkpoint directory",
+    )
+
+    return p.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    train_bandit(args)

train_rl.py

@@ -0,0 +1,674 @@
+#!/usr/bin/env python3
+"""REINFORCE with Learned Baseline -- RL training for PolypharmacyEnv.
+
+Trains a small neural-network policy to perform medication reviews in the
+PolypharmacyEnv environment.  The policy learns to query drug-drug interactions,
+propose clinical interventions, and decide when to finalise the review.
+
+Usage examples:
+    python train_rl.py --task easy_screening --episodes 200
+    python train_rl.py --task budgeted_screening --episodes 500
+    python train_rl.py --task complex_tradeoff --episodes 1000
+    python train_rl.py --task easy_screening --episodes 300 --lr 5e-4 --batch-size 8
+
+Architecture:
+    - Fixed-size state encoding (16-dim global summary features)
+    - Fixed 166-dim action space with dynamic validity masking
+    - 3-layer MLP policy  (state -> logits over actions)
+    - 3-layer MLP value baseline (state -> scalar return estimate)
+    - REINFORCE gradient with advantage = (discounted return) - baseline
+    - Entropy bonus for sustained exploration
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+from itertools import combinations
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical
+
+# ---------------------------------------------------------------------------
+# Environment imports (direct, no HTTP)
+# ---------------------------------------------------------------------------
+_BACKEND_SRC = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "backend", "src"
+)
+sys.path.insert(0, _BACKEND_SRC)
+
+from polypharmacy_env.env_core import PolypharmacyEnv  # noqa: E402
+from polypharmacy_env.models import (  # noqa: E402
+    PolypharmacyAction,
+    PolypharmacyObservation,
+)
+from polypharmacy_env.config import TASK_CONFIGS, TaskConfig  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Constants -- action-space geometry
+# ---------------------------------------------------------------------------
+MAX_MEDS = 15  # upper bound across all task difficulties
+INTERVENTION_TYPES: List[str] = [
+    "stop",
+    "dose_reduce",
+    "substitute",
+    "add_monitoring",
+]
+N_INTERVENTION_TYPES = len(INTERVENTION_TYPES)
+
+# Pre-compute the mapping (med_position_i, med_position_j) -> flat action index
+# for all possible query_ddi pairs where i < j.
+_PAIR_INDEX: Dict[Tuple[int, int], int] = {}
+_idx = 0
+for _i in range(MAX_MEDS):
+    for _j in range(_i + 1, MAX_MEDS):
+        _PAIR_INDEX[(_i, _j)] = _idx
+        _idx += 1
+N_PAIRS = _idx  # C(15,2) = 105
+_REVERSE_PAIR: Dict[int, Tuple[int, int]] = {v: k for k, v in _PAIR_INDEX.items()}
+
+N_INTERVENTIONS = MAX_MEDS * N_INTERVENTION_TYPES  # 60
+FINISH_IDX = N_PAIRS + N_INTERVENTIONS              # 165
+N_ACTIONS = FINISH_IDX + 1                           # 166
+
+# State feature vector length (see encode_state)
+STATE_DIM = 16
+
+# ---------------------------------------------------------------------------
+# State encoding
+# ---------------------------------------------------------------------------
+
+def encode_state(obs: PolypharmacyObservation, task_cfg: TaskConfig) -> torch.Tensor:
+    """Encode the observation into a compact 16-dim feature vector.
+
+    All values are normalised to roughly [0, 1] to help gradient flow.
+    """
+    meds = obs.current_medications
+    n_meds = len(meds)
+
+    n_high_risk = sum(1 for m in meds if m.is_high_risk_elderly)
+    n_beers_any = sum(1 for m in meds if m.beers_flags)
+    n_beers_avoid = sum(
+        1 for m in meds if any("avoid" in f for f in m.beers_flags)
+    )
+
+    queries = obs.interaction_queries
+    n_queries = len(queries)
+    n_severe = sum(1 for q in queries if q.severity == "severe")
+    n_moderate = sum(1 for q in queries if q.severity == "moderate")
+    n_interventions = len(obs.interventions)
+
+    max_possible_pairs = max(n_meds * (n_meds - 1) // 2, 1)
+
+    # Drugs involved in any discovered severe DDI (among current meds)
+    current_ids = {m.drug_id for m in meds}
+    drugs_in_severe: Set[str] = set()
+    for q in queries:
+        if q.severity == "severe":
+            if q.drug_id_1 in current_ids:
+                drugs_in_severe.add(q.drug_id_1)
+            if q.drug_id_2 in current_ids:
+                drugs_in_severe.add(q.drug_id_2)
+
+    features = [
+        n_meds / MAX_MEDS,
+        n_high_risk / max(n_meds, 1),
+        n_beers_any / max(n_meds, 1),
+        n_beers_avoid / max(n_meds, 1),
+        obs.remaining_query_budget / max(task_cfg.query_budget, 1),
+        obs.remaining_intervention_budget / max(task_cfg.intervention_budget, 1),
+        n_queries / max(task_cfg.query_budget, 1),
+        n_severe / max(n_queries, 1),
+        n_moderate / max(n_queries, 1),
+        n_interventions / max(task_cfg.intervention_budget, 1),
+        obs.step_index / max(task_cfg.max_steps, 1),
+        n_queries / max_possible_pairs,  # fraction of pairs queried
+        float(obs.remaining_query_budget > 0),
+        float(obs.remaining_intervention_budget > 0),
+        len(drugs_in_severe) / max(n_meds, 1),  # how much of the regimen is "hot"
+        float(n_meds <= 2),  # very few meds left -- may be time to finish
+    ]
+    return torch.tensor(features, dtype=torch.float32)
+
+
+# ---------------------------------------------------------------------------
+# Action-space helpers
+# ---------------------------------------------------------------------------
+
+def get_action_mask(obs: PolypharmacyObservation) -> torch.Tensor:
+    """Return a bool tensor of shape (N_ACTIONS,).  True = action is valid."""
+    mask = torch.zeros(N_ACTIONS, dtype=torch.bool)
+    meds = obs.current_medications
+    n_meds = min(len(meds), MAX_MEDS)
+
+    # Already-queried drug-id pairs (order-invariant)
+    queried: Set[frozenset] = set()
+    for q in obs.interaction_queries:
+        queried.add(frozenset((q.drug_id_1, q.drug_id_2)))
+
+    # --- query_ddi actions ---
+    if obs.remaining_query_budget > 0 and n_meds >= 2:
+        for i in range(n_meds):
+            for j in range(i + 1, n_meds):
+                pair_key = frozenset((meds[i].drug_id, meds[j].drug_id))
+                if pair_key not in queried:
+                    mask[_PAIR_INDEX[(i, j)]] = True
+
+    # --- propose_intervention actions ---
+    if obs.remaining_intervention_budget > 0:
+        for i in range(n_meds):
+            for k in range(N_INTERVENTION_TYPES):
+                mask[N_PAIRS + i * N_INTERVENTION_TYPES + k] = True
+
+    # --- finish_review (always valid) ---
+    mask[FINISH_IDX] = True
+    return mask
+
+
+def action_idx_to_env_action(
+    idx: int,
+    meds: list,
+) -> PolypharmacyAction:
+    """Map a flat action index back to a concrete PolypharmacyAction."""
+    if idx == FINISH_IDX:
+        return PolypharmacyAction(action_type="finish_review")
+
+    if idx < N_PAIRS:
+        i, j = _REVERSE_PAIR[idx]
+        return PolypharmacyAction(
+            action_type="query_ddi",
+            drug_id_1=meds[i].drug_id,
+            drug_id_2=meds[j].drug_id,
+        )
+
+    # Otherwise it is an intervention action
+    rel = idx - N_PAIRS
+    med_idx = rel // N_INTERVENTION_TYPES
+    type_idx = rel % N_INTERVENTION_TYPES
+    return PolypharmacyAction(
+        action_type="propose_intervention",
+        target_drug_id=meds[med_idx].drug_id,
+        intervention_type=INTERVENTION_TYPES[type_idx],
+        rationale="rl_policy",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Neural-network modules
+# ---------------------------------------------------------------------------
+
+class PolicyNetwork(nn.Module):
+    """3-layer MLP that maps state features to action logits."""
+
+    def __init__(
+        self,
+        state_dim: int = STATE_DIM,
+        action_dim: int = N_ACTIONS,
+        hidden: int = 128,
+    ) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(state_dim, hidden)
+        self.fc2 = nn.Linear(hidden, hidden)
+        self.fc3 = nn.Linear(hidden, action_dim)
+
+    def forward(
+        self,
+        state: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> Categorical:
+        x = F.relu(self.fc1(state))
+        x = F.relu(self.fc2(x))
+        logits = self.fc3(x)
+        logits = logits.masked_fill(~mask, float("-inf"))
+        return Categorical(logits=logits)
+
+
+class ValueNetwork(nn.Module):
+    """3-layer MLP baseline that estimates the expected return from a state."""
+
+    def __init__(self, state_dim: int = STATE_DIM, hidden: int = 128) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(state_dim, hidden)
+        self.fc2 = nn.Linear(hidden, hidden // 2)
+        self.fc3 = nn.Linear(hidden // 2, 1)
+
+    def forward(self, state: torch.Tensor) -> torch.Tensor:
+        x = F.relu(self.fc1(state))
+        x = F.relu(self.fc2(x))
+        return self.fc3(x).squeeze(-1)
+
+
+# ---------------------------------------------------------------------------
+# Episode rollout
+# ---------------------------------------------------------------------------
+
+def run_episode(
+    env: PolypharmacyEnv,
+    task_id: str,
+    policy: PolicyNetwork,
+    value_net: ValueNetwork,
+    task_cfg: TaskConfig,
+    seed: Optional[int] = None,
+    greedy: bool = False,
+) -> Dict[str, Any]:
+    """Roll out one full episode, collecting the REINFORCE trajectory.
+
+    When *greedy* is True the policy acts deterministically (argmax) and
+    gradients are not recorded.  Used for evaluation.
+    """
+    obs = env.reset(task_id=task_id, seed=seed)
+
+    states: List[torch.Tensor] = []
+    actions: List[torch.Tensor] = []
+    log_probs: List[torch.Tensor] = []
+    rewards: List[float] = []
+    values: List[torch.Tensor] = []
+    entropies: List[torch.Tensor] = []
+
+    grader_score = 0.0
+
+    while not obs.done:
+        state = encode_state(obs, task_cfg)
+        mask = get_action_mask(obs)
+
+        # Safety: if somehow no action is valid, force finish
+        if not mask.any():
+            mask[FINISH_IDX] = True
+
+        if greedy:
+            with torch.no_grad():
+                dist = policy(state, mask)
+                action_idx = dist.probs.argmax()
+                value = value_net(state)
+        else:
+            with torch.no_grad():
+                value = value_net(state)
+            dist = policy(state, mask)
+            action_idx = dist.sample()
+
+        log_prob = dist.log_prob(action_idx)
+        entropy = dist.entropy()
+
+        states.append(state)
+        actions.append(action_idx)
+        log_probs.append(log_prob)
+        values.append(value)
+        entropies.append(entropy)
+
+        env_action = action_idx_to_env_action(
+            action_idx.item(), obs.current_medications
+        )
+        obs = env.step(env_action)
+
+        reward = float(obs.reward) if obs.reward is not None else 0.0
+        rewards.append(reward)
+
+        if obs.done:
+            grader_score = obs.metadata.get("grader_score", 0.0)
+
+    return {
+        "states": states,
+        "actions": actions,
+        "log_probs": log_probs,
+        "rewards": rewards,
+        "values": values,
+        "entropies": entropies,
+        "grader_score": grader_score,
+        "total_reward": sum(rewards),
+        "n_steps": len(rewards),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Return computation
+# ---------------------------------------------------------------------------
+
+def compute_returns(rewards: List[float], gamma: float = 0.99) -> torch.Tensor:
+    """Discounted cumulative returns (G_t) for each timestep."""
+    returns: List[float] = []
+    g = 0.0
+    for r in reversed(rewards):
+        g = r + gamma * g
+        returns.insert(0, g)
+    return torch.tensor(returns, dtype=torch.float32)
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+def train(args: argparse.Namespace) -> None:  # noqa: C901 (complex but linear)
+    task_id: str = args.task
+    n_episodes: int = args.episodes
+    lr: float = args.lr
+    gamma: float = args.gamma
+    entropy_coeff: float = args.entropy_coeff
+    batch_size: int = args.batch_size
+    hidden_dim: int = args.hidden_dim
+    print_every: int = args.print_every
+
+    task_cfg = TASK_CONFIGS[task_id]
+
+    # ---- Initialise env & networks ----------------------------------------
+    env = PolypharmacyEnv()
+    policy = PolicyNetwork(STATE_DIM, N_ACTIONS, hidden=hidden_dim)
+    value_net = ValueNetwork(STATE_DIM, hidden=hidden_dim)
+
+    policy_optim = torch.optim.Adam(policy.parameters(), lr=lr)
+    value_optim = torch.optim.Adam(value_net.parameters(), lr=lr * 3)
+
+    # ---- Book-keeping -----------------------------------------------------
+    ckpt_dir = Path(args.checkpoint_dir)
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    metrics_path = Path(args.metrics_file)
+
+    episode_rewards: List[float] = []
+    episode_grader_scores: List[float] = []
+    episode_steps: List[int] = []
+    episode_policy_losses: List[float] = []
+    episode_value_losses: List[float] = []
+
+    best_avg_score: float = -float("inf")
+
+    print("=" * 72)
+    print("REINFORCE Training -- PolypharmacyEnv")
+    print("=" * 72)
+    print(f"  task            : {task_id}")
+    print(f"  episodes        : {n_episodes}")
+    print(f"  batch_size      : {batch_size}")
+    print(f"  lr              : {lr}")
+    print(f"  gamma           : {gamma}")
+    print(f"  entropy_coeff   : {entropy_coeff}")
+    print(f"  hidden_dim      : {hidden_dim}")
+    print(f"  state_dim       : {STATE_DIM}")
+    print(f"  action_space    : {N_ACTIONS}")
+    print(f"  task budgets    : query={task_cfg.query_budget}  "
+          f"intervention={task_cfg.intervention_budget}  "
+          f"max_steps={task_cfg.max_steps}")
+    print(f"  checkpoint_dir  : {ckpt_dir}")
+    print(f"  metrics_file    : {metrics_path}")
+    print("=" * 72)
+    print()
+
+    t_start = time.time()
+
+    # ---- Main training loop -----------------------------------------------
+    # Accumulate a mini-batch of trajectories, then perform one gradient step.
+    batch_trajs: List[Dict[str, Any]] = []
+
+    for ep in range(1, n_episodes + 1):
+        traj = run_episode(env, task_id, policy, value_net, task_cfg, seed=ep)
+
+        episode_rewards.append(traj["total_reward"])
+        episode_grader_scores.append(traj["grader_score"])
+        episode_steps.append(traj["n_steps"])
+
+        if traj["n_steps"] == 0:
+            # Degenerate episode (should not happen); skip update
+            continue
+
+        batch_trajs.append(traj)
+
+        # ---- Gradient step every batch_size episodes ----------------------
+        if len(batch_trajs) >= batch_size:
+            # Aggregate losses across the batch
+            total_policy_loss = torch.tensor(0.0)
+            total_value_loss = torch.tensor(0.0)
+            total_entropy = torch.tensor(0.0)
+            total_steps = 0
+
+            for bt in batch_trajs:
+                returns = compute_returns(bt["rewards"], gamma)
+                old_values_t = torch.stack(bt["values"])  # detached, from rollout
+                log_probs_t = torch.stack(bt["log_probs"])
+                entropies_t = torch.stack(bt["entropies"])
+
+                # Advantages use the *detached* rollout values as baseline
+                advantages = returns - old_values_t.detach()
+                # Per-trajectory advantage normalisation (reduces variance)
+                if len(advantages) > 1:
+                    advantages = (advantages - advantages.mean()) / (
+                        advantages.std() + 1e-8
+                    )
+
+                # REINFORCE policy gradient (negative because we minimise)
+                total_policy_loss = total_policy_loss + (
+                    -(log_probs_t * advantages).sum()
+                )
+
+                # Recompute value predictions WITH gradients for the value loss
+                states_t = torch.stack(bt["states"])
+                fresh_values = value_net(states_t)
+                total_value_loss = total_value_loss + F.mse_loss(
+                    fresh_values, returns, reduction="sum"
+                )
+
+                # Entropy (we want to maximise -> subtract from loss)
+                total_entropy = total_entropy + entropies_t.sum()
+                total_steps += len(bt["rewards"])
+
+            # Normalise by total number of timesteps in the batch
+            total_policy_loss = total_policy_loss / total_steps
+            total_value_loss = total_value_loss / total_steps
+            total_entropy = total_entropy / total_steps
+
+            # Combined policy loss with entropy bonus
+            combined_policy_loss = total_policy_loss - entropy_coeff * total_entropy
+
+            policy_optim.zero_grad()
+            combined_policy_loss.backward()
+            nn.utils.clip_grad_norm_(policy.parameters(), max_norm=1.0)
+            policy_optim.step()
+
+            value_optim.zero_grad()
+            total_value_loss.backward()
+            nn.utils.clip_grad_norm_(value_net.parameters(), max_norm=1.0)
+            value_optim.step()
+
+            episode_policy_losses.append(total_policy_loss.item())
+            episode_value_losses.append(total_value_loss.item())
+
+            batch_trajs = []
+
+        # ---- Logging ------------------------------------------------------
+        if ep % print_every == 0 or ep == 1:
+            window = min(print_every, ep)
+            recent_r = episode_rewards[-window:]
+            recent_s = episode_grader_scores[-window:]
+            recent_st = episode_steps[-window:]
+            avg_r = sum(recent_r) / len(recent_r)
+            avg_s = sum(recent_s) / len(recent_s)
+            avg_st = sum(recent_st) / len(recent_st)
+            elapsed = time.time() - t_start
+            print(
+                f"[ep {ep:>4d}/{n_episodes}]  "
+                f"avg_reward={avg_r:+.4f}  "
+                f"avg_grader={avg_s:.4f}  "
+                f"avg_steps={avg_st:.1f}  "
+                f"elapsed={elapsed:.1f}s"
+            )
+
+            # Save best checkpoint based on rolling grader score
+            eval_window = min(30, ep)
+            rolling_score = sum(episode_grader_scores[-eval_window:]) / eval_window
+            if rolling_score > best_avg_score:
+                best_avg_score = rolling_score
+                _save_checkpoint(
+                    policy, value_net, policy_optim, value_optim,
+                    ep, best_avg_score, task_id,
+                    ckpt_dir / f"best_{task_id}.pt",
+                )
+
+    # ---- Final checkpoint -------------------------------------------------
+    _save_checkpoint(
+        policy, value_net, policy_optim, value_optim,
+        n_episodes, best_avg_score, task_id,
+        ckpt_dir / f"final_{task_id}.pt",
+    )
+
+    # ---- Save training metrics to JSON ------------------------------------
+    metrics = {
+        "task_id": task_id,
+        "n_episodes": n_episodes,
+        "hyperparameters": {
+            "lr": lr,
+            "gamma": gamma,
+            "entropy_coeff": entropy_coeff,
+            "batch_size": batch_size,
+            "hidden_dim": hidden_dim,
+            "state_dim": STATE_DIM,
+            "action_dim": N_ACTIONS,
+        },
+        "episode_rewards": episode_rewards,
+        "episode_grader_scores": episode_grader_scores,
+        "episode_steps": episode_steps,
+        "policy_losses": episode_policy_losses,
+        "value_losses": episode_value_losses,
+        "best_avg_grader_score": best_avg_score,
+        "total_training_time_s": time.time() - t_start,
+    }
+    metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(metrics_path, "w") as f:
+        json.dump(metrics, f, indent=2)
+    print(f"\nTraining metrics saved to {metrics_path}")
+
+    # ---- Post-training evaluation -----------------------------------------
+    n_eval = 20
+    print("\n" + "=" * 72)
+    print(f"Post-training evaluation ({n_eval} episodes each mode)")
+    print("=" * 72)
+
+    for mode, is_greedy in [("stochastic", False), ("greedy", True)]:
+        eval_rewards, eval_scores, eval_steps_list = [], [], []
+        for i in range(n_eval):
+            traj = run_episode(
+                env, task_id, policy, value_net, task_cfg,
+                seed=10_000 + i, greedy=is_greedy,
+            )
+            eval_rewards.append(traj["total_reward"])
+            eval_scores.append(traj["grader_score"])
+            eval_steps_list.append(traj["n_steps"])
+        avg_r = sum(eval_rewards) / len(eval_rewards)
+        avg_s = sum(eval_scores) / len(eval_scores)
+        avg_st = sum(eval_steps_list) / len(eval_steps_list)
+        print(
+            f"  [{mode:>10s}]  avg_reward={avg_r:+.4f}  "
+            f"avg_grader={avg_s:.4f}  avg_steps={avg_st:.1f}"
+        )
+        metrics[f"eval_{mode}_avg_reward"] = avg_r
+        metrics[f"eval_{mode}_avg_grader_score"] = avg_s
+        metrics[f"eval_{mode}_avg_steps"] = avg_st
+        metrics[f"eval_{mode}_rewards"] = eval_rewards
+        metrics[f"eval_{mode}_grader_scores"] = eval_scores
+
+    print(f"  best training rolling-avg grader: {best_avg_score:.4f}")
+    print()
+
+    with open(metrics_path, "w") as f:
+        json.dump(metrics, f, indent=2)
+
+    print("Done.")
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint I/O
+# ---------------------------------------------------------------------------
+
+def _save_checkpoint(
+    policy: PolicyNetwork,
+    value_net: ValueNetwork,
+    policy_optim: torch.optim.Optimizer,
+    value_optim: torch.optim.Optimizer,
+    episode: int,
+    best_score: float,
+    task_id: str,
+    path: Path,
+) -> None:
+    torch.save(
+        {
+            "episode": episode,
+            "best_avg_grader_score": best_score,
+            "task_id": task_id,
+            "policy_state_dict": policy.state_dict(),
+            "value_state_dict": value_net.state_dict(),
+            "policy_optim_state_dict": policy_optim.state_dict(),
+            "value_optim_state_dict": value_optim.state_dict(),
+            "state_dim": STATE_DIM,
+            "action_dim": N_ACTIONS,
+        },
+        path,
+    )
+
+
+def load_checkpoint(
+    path: Path,
+    hidden_dim: int = 128,
+) -> Tuple[PolicyNetwork, ValueNetwork]:
+    """Load a trained policy + value net from a checkpoint file."""
+    ckpt = torch.load(path, map_location="cpu")
+    policy = PolicyNetwork(
+        ckpt.get("state_dim", STATE_DIM),
+        ckpt.get("action_dim", N_ACTIONS),
+        hidden=hidden_dim,
+    )
+    value_net = ValueNetwork(ckpt.get("state_dim", STATE_DIM), hidden=hidden_dim)
+    policy.load_state_dict(ckpt["policy_state_dict"])
+    value_net.load_state_dict(ckpt["value_state_dict"])
+    return policy, value_net
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="REINFORCE training for PolypharmacyEnv",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument(
+        "--task",
+        type=str,
+        default="easy_screening",
+        choices=list(TASK_CONFIGS.keys()),
+        help="Task difficulty to train on",
+    )
+    p.add_argument("--episodes", type=int, default=200, help="Number of training episodes")
+    p.add_argument("--batch-size", type=int, default=5, help="Episodes per gradient update")
+    p.add_argument("--lr", type=float, default=3e-4, help="Learning rate for Adam")
+    p.add_argument("--gamma", type=float, default=0.99, help="Discount factor")
+    p.add_argument(
+        "--entropy-coeff", type=float, default=0.02,
+        help="Entropy bonus coefficient (higher = more exploration)",
+    )
+    p.add_argument("--hidden-dim", type=int, default=128, help="Hidden layer width")
+    p.add_argument("--print-every", type=int, default=10, help="Print interval (episodes)")
+    p.add_argument(
+        "--checkpoint-dir",
+        type=str,
+        default=os.path.join(_BACKEND_SRC, "polypharmacy_env", "checkpoints"),
+        help="Directory to save model checkpoints",
+    )
+    p.add_argument(
+        "--metrics-file",
+        type=str,
+        default="training_metrics.json",
+        help="Path for JSON training metrics",
+    )
+    return p.parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    args = parse_args()
+    train(args)

training_metrics.json

@@ -0,0 +1,221 @@
+{
+  "task_id": "easy_screening",
+  "n_episodes": 30,
+  "hyperparameters": {
+    "lr": 0.0003,
+    "gamma": 0.99,
+    "entropy_coeff": 0.02,
+    "batch_size": 5,
+    "hidden_dim": 128,
+    "state_dim": 16,
+    "action_dim": 166
+  },
+  "episode_rewards": [
+    0.47,
+    1.1073118279569893,
+    1.1486231884057971,
+    1.1336231884057972,
+    0.405,
+    0.395,
+    1.1296806853582555,
+    0.38,
+    0.7283823529411766,
+    0.0,
+    0.39,
+    -0.095,
+    1.137962962962963,
+    1.1951785714285714,
+    0.9053636363636364,
+    -0.01754088050314473,
+    -0.04,
+    -0.06,
+    0.0,
+    0.22666666666666668,
+    0.435,
+    0.45,
+    0.45,
+    0.37666666666666665,
+    0.435,
+    0.455,
+    0.5412162162162163,
+    0.33899999999999997,
+    0.3416666666666666,
+    0.42
+  ],
+  "episode_grader_scores": [
+    0.5,
+    0.8306451612903226,
+    0.8369565217391305,
+    0.8369565217391305,
+    0.5,
+    0.5,
+    0.9688473520249221,
+    0.5,
+    0.7058823529411765,
+    0.0,
+    0.5,
+    0.0,
+    0.837962962962963,
+    0.9776785714285714,
+    0.8863636363636364,
+    0.053459119496855306,
+    0.0,
+    0.0,
+    0.0,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5495495495495496,
+    0.5,
+    0.27499999999999997,
+    0.5
+  ],
+  "episode_steps": [
+    5,
+    4,
+    3,
+    4,
+    6,
+    6,
+    5,
+    7,
+    4,
+    1,
+    7,
+    8,
+    4,
+    3,
+    10,
+    6,
+    3,
+    7,
+    1,
+    4,
+    4,
+    3,
+    3,
+    5,
+    4,
+    6,
+    4,
+    7,
+    2,
+    5
+  ],
+  "policy_losses": [
+    0.28967130184173584,
+    0.05730011314153671,
+    -0.06924888491630554,
+    -0.28697478771209717,
+    -0.1783256083726883,
+    -0.12063005566596985
+  ],
+  "value_losses": [
+    0.39626142382621765,
+    0.24146510660648346,
+    0.29013994336128235,
+    0.06388193368911743,
+    0.02375689707696438,
+    0.02241377718746662
+  ],
+  "best_avg_grader_score": 0.6179287909734683,
+  "total_training_time_s": 0.13345718383789062,
+  "eval_stochastic_avg_reward": 0.5792066304974347,
+  "eval_stochastic_avg_grader_score": 0.5784816304974348,
+  "eval_stochastic_avg_steps": 5.6,
+  "eval_stochastic_rewards": [
+    1.0402956989247312,
+    0.485,
+    -0.11333333333333336,
+    0.455,
+    0.36250000000000004,
+    1.0112089552238808,
+    0.41,
+    0.21499999999999997,
+    1.1173118279569894,
+    0.883621495327103,
+    0.45,
+    1.1073118279569893,
+    0.8680882352941177,
+    0.405,
+    0.675031128404669,
+    0.45,
+    -0.11,
+    0.415,
+    0.32,
+    1.1370967741935485
+  ],
+  "eval_stochastic_grader_scores": [
+    0.8494623655913979,
+    0.5,
+    0.0,
+    0.5,
+    0.5,
+    0.9207089552238806,
+    0.5,
+    0.5,
+    0.8306451612903226,
+    0.8411214953271029,
+    0.5,
+    0.8306451612903226,
+    0.803921568627451,
+    0.5,
+    0.6060311284046691,
+    0.5,
+    0.0,
+    0.5,
+    0.5,
+    0.8870967741935485
+  ],
+  "eval_greedy_avg_reward": 0.3627500000000001,
+  "eval_greedy_avg_grader_score": 0.425,
+  "eval_greedy_avg_steps": 6.45,
+  "eval_greedy_rewards": [
+    0.44,
+    0.455,
+    -0.08,
+    0.455,
+    -0.06,
+    0.39,
+    0.455,
+    0.455,
+    0.455,
+    -0.06,
+    0.42,
+    0.455,
+    0.41000000000000003,
+    0.455,
+    0.44,
+    0.39,
+    0.41000000000000003,
+    0.49,
+    0.44,
+    0.44
+  ],
+  "eval_greedy_grader_scores": [
+    0.5,
+    0.5,
+    0.0,
+    0.5,
+    0.0,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.0,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5,
+    0.5
+  ]
+}