Use deterministic task graders for all three tasks

graders.py

@@ -1,40 +1,273 @@
 """
-Static grader registry for OpenEnv task discoverability.
+Deterministic agent graders for the IT Mental Health OpenEnv tasks.
 
-The environment computes the real rubric score in `it_mental_health_environment.py`.
-These helpers expose one named grader per task so submission validators can
-statically discover at least three task-grader pairs.
+Each task has:
+- a concrete objective
+- a deterministic programmatic grader
+- a normalized score in the range [0.0, 1.0]
+
+These graders are intentionally rule-based so submission validators can
+statically discover them and evaluation remains reproducible.
 """
 
-from typing import Any, Dict
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, List, Tuple
+
+
+TaskState = Dict[str, Any]
+GradeResult = Tuple[float, Dict[str, float], str]
 
 
 def _normalize_reward(reward: float) -> float:
-    return min(max(float(reward), 0.0), 1.0)
+    return round(min(max(float(reward), 0.0), 1.0), 3)
 
 
-def _task_matches(state: Dict[str, Any], expected_task_id: str) -> bool:
-    if not isinstance(state, dict):
-        return False
+def _safe_lower(value: Any) -> str:
+    return str(value or "").lower()
 
-    candidates = [
-        state.get("task_id"),
-        state.get("current_task"),
-        (state.get("metadata") or {}).get("task_id") if isinstance(state.get("metadata"), dict) else None,
-    ]
+
+def _contains_any(text: str, terms: Iterable[str]) -> bool:
+    return any(term in text for term in terms)
+
+
+def _count_contains(text: str, terms: Iterable[str]) -> int:
+    return sum(1 for term in terms if term in text)
+
+
+def _metadata_task_id(state: TaskState) -> str:
+    metadata = state.get("metadata")
+    if isinstance(metadata, dict):
+        return str(metadata.get("task_id", ""))
+    return ""
+
+
+def _task_matches(state: TaskState, expected_task_id: str) -> bool:
+    candidates = {
+        str(state.get("task_id", "")),
+        str(state.get("current_task", "")),
+        _metadata_task_id(state),
+    }
     return expected_task_id in candidates
 
 
-def grade_burnout_detection(state: Dict[str, Any], reward: float) -> float:
-    return _normalize_reward(reward if _task_matches(state, "burnout_detection") else 0.0)
+def _breakdown_to_reward(
+    breakdown: Dict[str, float], weights: Dict[str, float]
+) -> float:
+    total_weight = sum(weights.values()) or 1.0
+    weighted = sum(breakdown.get(key, 0.0) * weight for key, weight in weights.items())
+    return _normalize_reward(weighted / total_weight)
+
+
+def _burnout_grade(response: str, ground_truth: Dict[str, Any]) -> GradeResult:
+    text = _safe_lower(response)
+    dims = ground_truth.get("active_dimensions", [])
+    dim_tokens = {
+        "Exhaustion": ["exhaustion", "exhausted", "fatigue", "tiredness"],
+        "Depersonalization": ["depersonalization", "depersonalisation", "cynicism", "cynical", "detached"],
+        "Reduced Personal Accomplishment": [
+            "reduced personal accomplishment",
+            "reduced accomplishment",
+            "inefficacy",
+            "ineffective",
+            "personal accomplishment",
+        ],
+    }
+
+    dim_hits = 0
+    for dimension in dims:
+        if _contains_any(text, dim_tokens.get(dimension, [dimension.lower()])):
+            dim_hits += 1
+    dimensions_identified = dim_hits / max(len(dims), 1)
+
+    severity_order = ["low", "moderate", "high", "critical"]
+    expected_severity = _safe_lower(ground_truth.get("severity"))
+    mentioned = [level for level in severity_order if level in text]
+    if expected_severity in mentioned:
+        severity_accuracy = 1.0
+    elif mentioned and expected_severity in severity_order:
+        gt_idx = severity_order.index(expected_severity)
+        min_distance = min(abs(severity_order.index(level) - gt_idx) for level in mentioned)
+        severity_accuracy = 0.5 if min_distance == 1 else 0.0
+    else:
+        severity_accuracy = 0.0
+
+    red_flag_terms = [
+        str(ground_truth.get("hours", "")),
+        str(ground_truth.get("vacation_months", "")),
+        "week",
+        "hours",
+        "vacation",
+        "leave",
+        "headache",
+        "migraine",
+        "sleep",
+        "desk",
+        "detached",
+        "cynical",
+        "incomplete",
+        "productivity",
+    ]
+    red_flags_quality = min(1.0, _count_contains(text, red_flag_terms) / 3.0)
+
+    escalation_expected = bool(ground_truth.get("escalation_needed"))
+    escalation_positive = _contains_any(text, ["yes", "escalat", "immediate hr", "urgent hr"])
+    escalation_reasoning = 1.0 if escalation_positive == escalation_expected else 0.0
+
+    structure_clarity = 1.0 if _contains_any(response, ["1.", "2.", "3.", "4.", "##", "**"]) else 0.4
+
+    breakdown = {
+        "dimensions_identified": round(dimensions_identified, 3),
+        "severity_accuracy": round(severity_accuracy, 3),
+        "red_flags_quality": round(red_flags_quality, 3),
+        "escalation_reasoning": round(escalation_reasoning, 3),
+        "structure_clarity": round(structure_clarity, 3),
+    }
+    weights = {
+        "dimensions_identified": 0.30,
+        "severity_accuracy": 0.20,
+        "red_flags_quality": 0.20,
+        "escalation_reasoning": 0.20,
+        "structure_clarity": 0.10,
+    }
+    reward = _breakdown_to_reward(breakdown, weights)
+    feedback = (
+        f"[Deterministic Grader] burnout_detection score={reward:.2f}. "
+        f"Expected severity={ground_truth.get('severity')}; "
+        f"dimension hits={dim_hits}/{max(len(dims), 1)}."
+    )
+    return reward, breakdown, feedback
+
+
+def _stress_triage_grade(response: str, ground_truth: Dict[str, Any]) -> GradeResult:
+    text = _safe_lower(response)
+    tiers = ground_truth.get("correct_tiers", {})
+    names = ground_truth.get("names", [])
+    priority_order = [str(name).lower() for name in ground_truth.get("priority_order", [])]
+
+    tier_hits = 0
+    for name, tier in tiers.items():
+        name_lower = name.lower()
+        tier_lower = tier.lower()
+        if name_lower in text and tier_lower in text:
+            tier_hits += 1
+    tier_accuracy = tier_hits / max(len(tiers), 1)
+
+    top_two_correct = sum(1 for name in priority_order[:2] if name and name in text[:800])
+    priority_ranking = 1.0 if top_two_correct == 2 else 0.5 if top_two_correct == 1 else 0.0
 
+    immediate_actions = 1.0 if _contains_any(text, ["24 hour", "24-hour", "within 24", "immediate", "today"]) else 0.3
+    medium_term_support = 1.0 if _contains_any(text, ["2 week", "2-week", "within 2 weeks", "fortnight", "two weeks"]) else 0.3
+    completeness = 1.0 if all(name.lower() in text for name in names) else 0.0
 
-def grade_stress_triage(state: Dict[str, Any], reward: float) -> float:
-    return _normalize_reward(reward if _task_matches(state, "stress_triage") else 0.0)
+    breakdown = {
+        "tier_accuracy": round(tier_accuracy, 3),
+        "priority_ranking": round(priority_ranking, 3),
+        "immediate_actions": round(immediate_actions, 3),
+        "medium_term_support": round(medium_term_support, 3),
+        "completeness": round(completeness, 3),
+    }
+    weights = {
+        "tier_accuracy": 0.35,
+        "priority_ranking": 0.20,
+        "immediate_actions": 0.15,
+        "medium_term_support": 0.15,
+        "completeness": 0.15,
+    }
+    reward = _breakdown_to_reward(breakdown, weights)
+    feedback = (
+        f"[Deterministic Grader] stress_triage score={reward:.2f}. "
+        f"Correct tier assignments={tier_hits}/{max(len(tiers), 1)}."
+    )
+    return reward, breakdown, feedback
 
 
-def grade_intervention_plan(state: Dict[str, Any], reward: float) -> float:
-    return _normalize_reward(reward if _task_matches(state, "intervention_plan") else 0.0)
+def _intervention_plan_grade(response: str, ground_truth: Dict[str, Any]) -> GradeResult:
+    text = _safe_lower(response)
+
+    week_hits = sum(1 for week in ["week 1", "week 2", "week 3", "week 4"] if week in text)
+    four_week_structure = week_hits / 4.0
+
+    action_terms = [
+        "on-call",
+        "overtime",
+        "survey",
+        "1:1",
+        "1-1",
+        "check-in",
+        "training",
+        "leave",
+        "workload",
+        "rotation",
+        "policy",
+    ]
+    action_concreteness = min(1.0, _count_contains(text, action_terms) / 6.0)
+    responsibility = 1.0 if _contains_any(text, ["hr", "manager", "eap"]) else 0.0
+
+    kpi_terms = ["kpi", "metric", "measure", "indicator", "tracking", "90-day"]
+    kpis_quality = 1.0 if _count_contains(text, kpi_terms) >= 2 else 0.5 if _count_contains(text, kpi_terms) == 1 else 0.0
+
+    risk_and_budget = 1.0 if ("risk" in text and _contains_any(text, ["budget", "low", "medium", "high", "$"])) else 0.0
+
+    proportionality_terms = [
+        str(ground_truth.get("team_size", "")),
+        str(ground_truth.get("affected", "")),
+        str(ground_truth.get("overtime", "")),
+        str(ground_truth.get("oncall_days", "")),
+        str(ground_truth.get("hr_complaints", "")),
+    ]
+    proportionality = 1.0 if _count_contains(text, [term for term in proportionality_terms if term]) >= 2 else 0.5
+
+    breakdown = {
+        "four_week_structure": round(four_week_structure, 3),
+        "action_concreteness": round(action_concreteness, 3),
+        "responsibility": round(responsibility, 3),
+        "kpis_quality": round(kpis_quality, 3),
+        "risk_and_budget": round(risk_and_budget, 3),
+        "proportionality": round(proportionality, 3),
+    }
+    weights = {
+        "four_week_structure": 0.30,
+        "action_concreteness": 0.20,
+        "responsibility": 0.10,
+        "kpis_quality": 0.15,
+        "risk_and_budget": 0.10,
+        "proportionality": 0.15,
+    }
+    reward = _breakdown_to_reward(breakdown, weights)
+    feedback = (
+        f"[Deterministic Grader] intervention_plan score={reward:.2f}. "
+        f"Week coverage={week_hits}/4."
+    )
+    return reward, breakdown, feedback
+
+
+def grade_response(task_id: str, response: str, ground_truth: Dict[str, Any]) -> GradeResult:
+    if task_id == "burnout_detection":
+        return _burnout_grade(response, ground_truth)
+    if task_id == "stress_triage":
+        return _stress_triage_grade(response, ground_truth)
+    if task_id == "intervention_plan":
+        return _intervention_plan_grade(response, ground_truth)
+    return 0.0, {}, f"[Deterministic Grader] unknown task_id={task_id}"
+
+
+def grade_burnout_detection(state: TaskState, reward: float) -> float:
+    if not _task_matches(state, "burnout_detection"):
+        return 0.0
+    return _normalize_reward(reward)
+
+
+def grade_stress_triage(state: TaskState, reward: float) -> float:
+    if not _task_matches(state, "stress_triage"):
+        return 0.0
+    return _normalize_reward(reward)
+
+
+def grade_intervention_plan(state: TaskState, reward: float) -> float:
+    if not _task_matches(state, "intervention_plan"):
+        return 0.0
+    return _normalize_reward(reward)
 
 
 GRADERS = {
@@ -44,17 +277,26 @@ GRADERS = {
 }
 
 
-TASK_GRADER_PAIRS = [
+TASK_GRADER_PAIRS: List[Tuple[str, str]] = [
     ("burnout_detection", "burnout_detection_grader"),
     ("stress_triage", "stress_triage_grader"),
     ("intervention_plan", "intervention_plan_grader"),
 ]
 
 
+TASK_GRADER_OBJECTIVES = {
+    "burnout_detection": "Identify burnout dimensions, severity, red flags, and escalation need from an employee profile.",
+    "stress_triage": "Triage three employees by urgency and recommend immediate plus medium-term support.",
+    "intervention_plan": "Produce a four-week intervention plan with owners, KPIs, risk, and budget.",
+}
+
+
 __all__ = [
-    "grade_burnout_detection",
-    "grade_stress_triage",
-    "grade_intervention_plan",
     "GRADERS",
+    "TASK_GRADER_OBJECTIVES",
     "TASK_GRADER_PAIRS",
+    "grade_burnout_detection",
+    "grade_intervention_plan",
+    "grade_response",
+    "grade_stress_triage",
 ]

it_mental_health_environment.py

@@ -1,26 +1,29 @@
 """
-IT Mental Health OpenEnv - Environment Logic  (v2 - improved)
+IT Mental Health OpenEnv - Environment Logic.
 
-FIX 1: LLM-as-judge grader replaces keyword matching.
-FIX 2: Randomised episode generation — every reset() is unique.
+This environment defines three tasks with deterministic programmatic graders:
+- burnout_detection (easy)
+- stress_triage (medium)
+- intervention_plan (hard)
 """
 
-import os
-import uuid
 import random
-import json
-from typing import Optional
-from openai import OpenAI
+import uuid
+
+from graders import grade_response
 
 try:
-    from models import MentalHealthAction, MentalHealthObservation, MentalHealthReward, MentalHealthState
+    from models import MentalHealthObservation, MentalHealthReward, MentalHealthState
 except ImportError:
     try:
-        from server.models import MentalHealthAction, MentalHealthObservation, MentalHealthReward, MentalHealthState
+        from server.models import MentalHealthObservation, MentalHealthReward, MentalHealthState
     except ImportError:
+        import os
         import sys
+
         sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-        from models import MentalHealthAction, MentalHealthObservation, MentalHealthReward, MentalHealthState
+        from models import MentalHealthObservation, MentalHealthReward, MentalHealthState
+
 
 TASK_ORDER = ["burnout_detection", "stress_triage", "intervention_plan"]
 TASK_DIFFICULTY = {
@@ -29,27 +32,19 @@ TASK_DIFFICULTY = {
     "intervention_plan": "hard",
 }
 
-# ── LLM Judge client ──────────────────────────────────────────────────────────
-_llm_client: Optional[OpenAI] = None
-
-def _get_llm_client() -> Optional[OpenAI]:
-    global _llm_client
-    use_llm_judge = os.environ.get("USE_LLM_JUDGE", "").strip().lower() in {"1", "true", "yes"}
-    if not use_llm_judge:
-        return None
-    if _llm_client is None:
-        api_key = os.environ.get("HF_TOKEN", "") or os.environ.get("OPENAI_API_KEY", "")
-        base_url = os.environ.get("API_BASE_URL", "https://api-inference.huggingface.co/v1")
-        if api_key:
-            _llm_client = OpenAI(api_key=api_key, base_url=base_url)
-    return _llm_client
 
-
-# ── Randomised scenario data pools ───────────────────────────────────────────
-NAMES = ["Alex","Jordan","Sam","Riley","Morgan","Casey","Taylor","Drew","Jamie","Avery"]
-ROLES = ["Software Engineer","DevOps Engineer","QA Engineer","Data Scientist",
-         "Backend Developer","Frontend Developer","ML Engineer","Platform Engineer"]
-YEARS_EXP = [1,2,3,4,5,7,8,10]
+NAMES = ["Alex", "Jordan", "Sam", "Riley", "Morgan", "Casey", "Taylor", "Drew", "Jamie", "Avery"]
+ROLES = [
+    "Software Engineer",
+    "DevOps Engineer",
+    "QA Engineer",
+    "Data Scientist",
+    "Backend Developer",
+    "Frontend Developer",
+    "ML Engineer",
+    "Platform Engineer",
+]
+YEARS_EXP = [1, 2, 3, 4, 5, 7, 8, 10]
 
 EXHAUSTION_SYMPTOMS = [
     "sleeping only 4-5 hours/night",
@@ -78,21 +73,21 @@ REDUCED_ACCOMPLISHMENT_SYMPTOMS = [
 ]
 
 SEVERITY_MAP = {
-    "Low":      {"hours_range":(42,50), "vacation_months":(1,4),  "dimensions":1},
-    "Moderate": {"hours_range":(50,60), "vacation_months":(4,8),  "dimensions":2},
-    "High":     {"hours_range":(60,68), "vacation_months":(8,14), "dimensions":2},
-    "Critical": {"hours_range":(68,80), "vacation_months":(14,24),"dimensions":3},
+    "Low": {"hours_range": (42, 50), "vacation_months": (1, 4), "dimensions": 1},
+    "Moderate": {"hours_range": (50, 60), "vacation_months": (4, 8), "dimensions": 2},
+    "High": {"hours_range": (60, 68), "vacation_months": (8, 14), "dimensions": 2},
+    "Critical": {"hours_range": (68, 80), "vacation_months": (14, 24), "dimensions": 3},
 }
 
 STRESS_TIER_TEMPLATES = {
-    "GREEN":    ["Feeling slightly overwhelmed with the sprint but I think I'll manage."],
-    "AMBER":    ["Haven't slept great — around 6 hours most nights. Starting to affect focus."],
-    "RED":      ["I snapped at a teammate yesterday. I'm scared of how I'm feeling lately."],
+    "GREEN": ["Feeling slightly overwhelmed with the sprint but I think I'll manage."],
+    "AMBER": ["Haven't slept great - around 6 hours most nights. Starting to affect focus."],
+    "RED": ["I snapped at a teammate yesterday. I'm scared of how I'm feeling lately."],
     "CRITICAL": ["I've been having chest tightness every time I get a pager alert. Three weeks. Haven't told anyone."],
 }
 
 
-def generate_burnout_scenario(rng):
+def generate_burnout_scenario(rng: random.Random):
     name = rng.choice(NAMES)
     role = rng.choice(ROLES)
     exp = rng.choice(YEARS_EXP)
@@ -101,7 +96,7 @@ def generate_burnout_scenario(rng):
     vacation = rng.randint(*cfg["vacation_months"])
     n_dims = cfg["dimensions"]
 
-    all_dims = ["exhaustion","depersonalization","reduced_accomplishment"]
+    all_dims = ["exhaustion", "depersonalization", "reduced_accomplishment"]
     active_dims = rng.sample(all_dims, n_dims)
     symptoms = []
     ground_truth_dims = []
@@ -115,9 +110,9 @@ def generate_burnout_scenario(rng):
         symptoms += rng.sample(REDUCED_ACCOMPLISHMENT_SYMPTOMS, 2)
         ground_truth_dims.append("Reduced Personal Accomplishment")
     rng.shuffle(symptoms)
-    sym_text = "\n".join(f"- {s}" for s in symptoms)
+    sym_text = "\n".join(f"- {item}" for item in symptoms)
 
-    scenario = f"""[TASK: Burnout Detection — EASY]
+    scenario = f"""[TASK: Burnout Detection - EASY]
 
 You are an occupational psychologist reviewing an IT employee profile.
 
@@ -128,74 +123,81 @@ Employee Profile:
 - Observed symptoms:
 {sym_text}
 
-Your task:
-1. Identify which MBI dimensions are present: Exhaustion / Depersonalization / Reduced Personal Accomplishment
+Objective:
+1. Identify which MBI dimensions are present
 2. Rate severity: Low / Moderate / High / Critical
-3. List TOP 3 red-flag signals from the profile
-4. State whether immediate HR escalation is needed (Yes/No) and why
+3. List the top 3 red-flag signals
+4. State whether immediate HR escalation is needed and why
 
 Respond with clear headings."""
 
-    gt = {"active_dimensions":ground_truth_dims, "severity":severity_label,
-          "escalation_needed":severity_label in ("High","Critical"),
-          "name":name, "hours":hours, "vacation_months":vacation}
-    return scenario, gt
+    ground_truth = {
+        "active_dimensions": ground_truth_dims,
+        "severity": severity_label,
+        "escalation_needed": severity_label in ("High", "Critical"),
+        "name": name,
+        "hours": hours,
+        "vacation_months": vacation,
+    }
+    return scenario, ground_truth
 
 
-def generate_stress_triage_scenario(rng):
-    tiers = ["CRITICAL","RED",rng.choice(["GREEN","AMBER"])]
+def generate_stress_triage_scenario(rng: random.Random):
+    tiers = ["CRITICAL", "RED", rng.choice(["GREEN", "AMBER"])]
     rng.shuffle(tiers)
     selected_names = rng.sample(NAMES, len(tiers))
     employees = []
     for idx, tier in enumerate(tiers):
-        employees.append({
-            "name": selected_names[idx],
-            "role": rng.choice(ROLES),
-            "tier": tier,
-            "quote": rng.choice(STRESS_TIER_TEMPLATES[tier]),
-        })
-    tier_order = {"CRITICAL":0,"RED":1,"AMBER":2,"GREEN":3}
-    priority_order = sorted(range(3), key=lambda i: tier_order[employees[i]["tier"]])
+        employees.append(
+            {
+                "name": selected_names[idx],
+                "role": rng.choice(ROLES),
+                "tier": tier,
+                "quote": rng.choice(STRESS_TIER_TEMPLATES[tier]),
+            }
+        )
 
+    tier_order = {"CRITICAL": 0, "RED": 1, "AMBER": 2, "GREEN": 3}
+    priority_order = sorted(range(3), key=lambda index: tier_order[employees[index]["tier"]])
     cases = "".join(
-        f'{i+1}. {e["name"]} ({e["role"]}): "{e["quote"]}"\n'
-        for i,e in enumerate(employees)
+        f'{index + 1}. {employee["name"]} ({employee["role"]}): "{employee["quote"]}"\n'
+        for index, employee in enumerate(employees)
     )
-    scenario = f"""[TASK: Stress Triage — MEDIUM]
+    scenario = f"""[TASK: Stress Triage - MEDIUM]
 
 You are reviewing urgent mental health flags from an IT pulse survey.
 Triage these 3 employees and prioritise them.
 
 Cases:
 {cases}
-For each case:
+Objective:
 a) Assign: GREEN / AMBER / RED / CRITICAL
-b) Primary stressor: Workload / Physiological / Relationship / Cognitive
-c) ONE immediate action (within 24 hours)
-d) ONE medium-term support (within 2 weeks)
+b) Name the primary stressor
+c) Give one immediate action within 24 hours
+d) Give one medium-term support action within 2 weeks
 
 Rank the 3 cases by intervention priority (1 = most urgent)."""
 
-    gt = {
+    ground_truth = {
         "employees": employees,
-        "correct_tiers": {e["name"]:e["tier"] for e in employees},
-        "priority_order": [employees[i]["name"] for i in priority_order],
-        "names": [e["name"] for e in employees],
+        "correct_tiers": {employee["name"]: employee["tier"] for employee in employees},
+        "priority_order": [employees[index]["name"] for index in priority_order],
+        "names": [employee["name"] for employee in employees],
     }
-    return scenario, gt
+    return scenario, ground_truth
 
 
-def generate_intervention_plan_scenario(rng):
-    team_size = rng.choice([8,10,12,15,18,20])
-    affected = int(team_size * rng.choice([0.4,0.5,0.6,0.7]))
-    overtime = rng.choice([10,15,18,20,25])
-    oncall = rng.choice([2,3,4,5])
-    no_tb = rng.choice([6,9,12,15,18,24])
-    no_1on1 = rng.choice([3,4,5,6,8])
-    on_leave = rng.randint(1,3)
-    hr_c = rng.randint(1,4)
+def generate_intervention_plan_scenario(rng: random.Random):
+    team_size = rng.choice([8, 10, 12, 15, 18, 20])
+    affected = int(team_size * rng.choice([0.4, 0.5, 0.6, 0.7]))
+    overtime = rng.choice([10, 15, 18, 20, 25])
+    oncall = rng.choice([2, 3, 4, 5])
+    no_tb = rng.choice([6, 9, 12, 15, 18, 24])
+    no_1on1 = rng.choice([3, 4, 5, 6, 8])
+    on_leave = rng.randint(1, 3)
+    hr_complaints = rng.randint(1, 4)
 
-    scenario = f"""[TASK: Intervention Plan — HARD]
+    scenario = f"""[TASK: Intervention Plan - HARD]
 
 You are a workplace mental health consultant. Design a 4-week intervention for a
 software team with systemic burnout ({affected} of {team_size} members affected).
@@ -205,167 +207,32 @@ Team Context:
 - On-call: 1 person every {oncall} days (24/7)
 - No team-building in {no_tb} months
 - No manager 1:1s in {no_1on1} months
-- {hr_c} formal HR complaints about workload
+- {hr_complaints} formal HR complaints about workload
 - {on_leave} members on anxiety-related medical leave
 
-Design a 4-week plan:
-Week 1: Immediate stabilisation
-Week 2: Assessment & listening
-Week 3: Process reforms (on-call, overtime, workload)
-Week 4: Sustainable systems
-
-For each week: 2+ concrete actions, responsible party (HR/Manager/EAP), measurable outcome.
-
-Also include:
-- 3 measurable KPIs for 90-day tracking
-- 1 key risk if plan is NOT executed
-- Budget: Low (<$500) / Medium ($500-$5000) / High (>$5000)"""
-
-    gt = {"team_size":team_size,"affected":affected,"overtime":overtime,
-          "oncall_days":oncall,"on_leave":on_leave,"hr_complaints":hr_c}
-    return scenario, gt
-
-
-# ── Clinical rubrics ──────────────────────────────────────────────────────────
-RUBRICS = {
-    "burnout_detection": {
-        "dimensions_identified": {"max":3.0,"desc":"1pt per correct MBI dimension identified (0-3)"},
-        "severity_accuracy":     {"max":2.0,"desc":"2=exact match, 1=adjacent, 0=wrong"},
-        "red_flags_quality":     {"max":2.0,"desc":"Flags are specific and drawn from the profile (0-2)"},
-        "escalation_reasoning":  {"max":2.0,"desc":"Correct Yes/No + sound clinical reasoning (0-2)"},
-        "structure_clarity":     {"max":1.0,"desc":"Clear headings, professional format (0-1)"},
-    },
-    "stress_triage": {
-        "tier_accuracy":       {"max":3.0,"desc":"1pt per correct tier assigned (0-3)"},
-        "priority_ranking":    {"max":2.0,"desc":"2=correct, 1=partially correct, 0=wrong"},
-        "immediate_actions":   {"max":2.0,"desc":"Specific 24h actions proportionate to tier (0-2)"},
-        "medium_term_support": {"max":2.0,"desc":"Realistic 2-week support for each case (0-2)"},
-        "completeness":        {"max":1.0,"desc":"All 3 employees addressed with all 4 elements (0-1)"},
-    },
-    "intervention_plan": {
-        "four_week_structure":  {"max":3.0,"desc":"All 4 weeks with distinct focus (0-3)"},
-        "action_concreteness":  {"max":2.0,"desc":"Actions specific, not vague (0-2)"},
-        "responsibility":       {"max":1.0,"desc":"Responsible party per action (0-1)"},
-        "kpis_quality":         {"max":2.0,"desc":"3 measurable, relevant KPIs (0-2)"},
-        "risk_and_budget":      {"max":1.0,"desc":"Risk realistic + budget category present (0-1)"},
-        "proportionality":      {"max":1.0,"desc":"Plan addresses this team's specific numbers (0-1)"},
-    },
-}
-
-JUDGE_SYSTEM = """You are a clinical rubric evaluator for an AI mental health RL environment.
-Score the agent's response against the rubric. Award marks based on reasoning quality and accuracy — NOT keyword presence alone.
-Return ONLY valid JSON, no markdown fences:
-{"scores": {"dim_name": float, ...}, "feedback": "2-3 sentence critique"}"""
-
-
-def _llm_judge_grade(task_id, scenario, response, ground_truth):
-    rubric = RUBRICS[task_id]
-    total_max = sum(v["max"] for v in rubric.values())
-    rubric_text = "\n".join(f'  "{k}": max={v["max"]} — {v["desc"]}' for k,v in rubric.items())
-
-    prompt = f"""Task: {task_id}
-Scenario: {scenario}
-Ground truth: {json.dumps(ground_truth)}
-Agent response: {response}
-Rubric:
-{rubric_text}
-Score now. JSON only."""
+Objective:
+- Week 1: Immediate stabilisation
+- Week 2: Assessment and listening
+- Week 3: Process reforms
+- Week 4: Sustainable systems
+
+For each week include 2+ actions, responsible owner, measurable outcome.
+Also include 3 KPIs, 1 key risk, and a budget category."""
+
+    ground_truth = {
+        "team_size": team_size,
+        "affected": affected,
+        "overtime": overtime,
+        "oncall_days": oncall,
+        "on_leave": on_leave,
+        "hr_complaints": hr_complaints,
+    }
+    return scenario, ground_truth
 
-    client = _get_llm_client()
-    if client is None:
-        return _heuristic_grade(task_id, response, ground_truth)
 
-    model = os.environ.get("MODEL_NAME","meta-llama/Llama-3.1-8B-Instruct")
-    try:
-        resp = client.chat.completions.create(
-            model=model,
-            messages=[{"role":"system","content":JUDGE_SYSTEM},
-                      {"role":"user","content":prompt}],
-            max_tokens=400, temperature=0.0,
-        )
-        raw = resp.choices[0].message.content.strip().replace("```json","").replace("```","").strip()
-        parsed = json.loads(raw)
-        scores = parsed.get("scores",{})
-        feedback = parsed.get("feedback","")
-        breakdown, total = {}, 0.0
-        for dim, cfg in rubric.items():
-            s = max(0.0, min(cfg["max"], float(scores.get(dim, 0.0))))
-            breakdown[dim] = round(s/cfg["max"], 3)
-            total += s
-        reward = round(max(0.0, min(1.0, total/total_max)), 3)
-        return reward, breakdown, f"[LLM Judge] {feedback}"
-    except Exception as e:
-        r,b,f = _heuristic_grade(task_id, response, ground_truth)
-        return r, b, f"[Fallback — judge error: {e}] {f}"
-
-
-def _heuristic_grade(task_id, response, ground_truth):
-    """Ground-truth-aware heuristic fallback (better than pure keyword matching)."""
-    r = response.lower()
-    breakdown, score = {}, 0.0
-
-    if task_id == "burnout_detection":
-        dims = ground_truth.get("active_dimensions",[])
-        # Match on a distinctive token per dimension, not just the first word.
-        dim_keys = {
-            "Exhaustion": ["exhaustion","exhausted"],
-            "Depersonalization": ["depersonalization","depersonalisation","cynicism","cynical"],
-            "Reduced Personal Accomplishment": ["reduced personal accomplishment","reduced accomplishment","inefficacy","personal accomplishment"],
-        }
-        dim_hits = sum(1 for d in dims if any(k in r for k in dim_keys.get(d, [d.lower()])))
-        breakdown["dimensions_identified"] = dim_hits/max(len(dims),1)
-        # Severity: exact match = 1.0, adjacent tier = 0.5, otherwise 0.
-        sev_order = ["low","moderate","high","critical"]
-        sev = ground_truth.get("severity","High").lower()
-        mentioned = [s for s in sev_order if s in r]
-        if sev in mentioned:
-            breakdown["severity_accuracy"] = 1.0
-        elif mentioned and sev in sev_order:
-            gt_idx = sev_order.index(sev)
-            best = min(abs(sev_order.index(m) - gt_idx) for m in mentioned)
-            breakdown["severity_accuracy"] = 0.5 if best == 1 else 0.0
-        else:
-            breakdown["severity_accuracy"] = 0.0
-        breakdown["red_flags_quality"] = min(1.0, len(response)/300)
-        esc = ground_truth.get("escalation_needed",True)
-        esc_hit = "yes" in r or "escalat" in r or "immediate" in r
-        breakdown["escalation_reasoning"] = 1.0 if esc==esc_hit else 0.3
-        breakdown["structure_clarity"] = 1.0 if any(h in response for h in ["##","**","1.","1:"]) else 0.3
-
-    elif task_id == "stress_triage":
-        tiers = ground_truth.get("correct_tiers",{})
-        tier_hits = sum(1 for nm,t in tiers.items() if nm.lower() in r and t.lower() in r)
-        breakdown["tier_accuracy"] = tier_hits/max(len(tiers),1)
-        priority = ground_truth.get("priority_order",[])
-        breakdown["priority_ranking"] = 1.0 if (priority and priority[0].lower() in r[:500]) else 0.4
-        breakdown["immediate_actions"] = min(1.0, r.count("24")*0.3 + r.count("immediat")*0.4)
-        breakdown["medium_term_support"] = 1.0 if ("week" in r or "fortnight" in r) else 0.3
-        names = ground_truth.get("names",[])
-        breakdown["completeness"] = 1.0 if all(n.lower() in r for n in names) else 0.4
-
-    elif task_id == "intervention_plan":
-        weeks = sum(1 for w in ["week 1","week 2","week 3","week 4"] if w in r)
-        breakdown["four_week_structure"] = weeks/4.0
-        breakdown["action_concreteness"] = min(1.0, len(response)/600)
-        breakdown["responsibility"] = 1.0 if any(x in r for x in ["hr","manager","eap"]) else 0.0
-        kpi_hits = sum(1 for k in ["kpi","metric","measure","indicator"] if k in r)
-        breakdown["kpis_quality"] = min(1.0, kpi_hits*0.4)
-        breakdown["risk_and_budget"] = 1.0 if ("risk" in r and any(b in r for b in ["low","medium","high","$"])) else 0.3
-        gt_size = str(ground_truth.get("team_size",""))
-        breakdown["proportionality"] = 1.0 if gt_size in response else 0.5
-
-    score = sum(breakdown.values())/max(len(breakdown),1)
-    reward = round(max(0.0, min(1.0, score)), 3)
-    return reward, breakdown, f"[Heuristic] {task_id} = {reward:.2f}. Breakdown: {breakdown}"
-
-
-# ── Environment class ─────────────────────────────────────────────────────────
 class ITMentalHealthEnvironment:
-    """
-    OpenEnv RL environment — IT sector mental health.
-    Every reset() is unique (randomised profiles/scenarios).
-    Graded by LLM-as-judge against clinical rubric.
-    """
+    """OpenEnv-compatible environment with deterministic task graders."""
+
     def __init__(self):
         self._state = MentalHealthState()
         self._task_index = 0
@@ -375,7 +242,7 @@ class ITMentalHealthEnvironment:
     def _generate_all(self):
         self._current_scenarios = {
             "burnout_detection": generate_burnout_scenario(self._rng),
-            "stress_triage":     generate_stress_triage_scenario(self._rng),
+            "stress_triage": generate_stress_triage_scenario(self._rng),
             "intervention_plan": generate_intervention_plan_scenario(self._rng),
         }
 
@@ -383,9 +250,7 @@ class ITMentalHealthEnvironment:
         difficulty = TASK_DIFFICULTY[task]
         task_step_count = self._state.task_step_counts.get(task, 0) + 1
         task_cumulative_score = round(self._state.task_scores.get(task, 0.0) + reward, 3)
-        difficulty_cumulative_score = round(
-            self._state.difficulty_scores.get(difficulty, 0.0) + reward, 3
-        )
+        difficulty_cumulative_score = round(self._state.difficulty_scores.get(difficulty, 0.0) + reward, 3)
         return {
             "difficulty": difficulty,
             "step_score": round(reward, 3),
@@ -393,6 +258,7 @@ class ITMentalHealthEnvironment:
             "task_cumulative_score": task_cumulative_score,
             "difficulty_cumulative_score": difficulty_cumulative_score,
             "overall_cumulative_score": round(self._state.cumulative_reward + reward, 3),
+            "task_id": task,
         }
 
     def reset(self, seed=None):
@@ -401,8 +267,11 @@ class ITMentalHealthEnvironment:
         self._task_index = 0
         self._generate_all()
         self._state = MentalHealthState(
-            episode_id=str(uuid.uuid4()), step_count=0,
-            current_task=TASK_ORDER[0], cumulative_reward=0.0, tasks_completed=[],
+            episode_id=str(uuid.uuid4()),
+            step_count=0,
+            current_task=TASK_ORDER[0],
+            cumulative_reward=0.0,
+            tasks_completed=[],
             task_scores={task_id: 0.0 for task_id in TASK_ORDER},
             task_step_counts={task_id: 0 for task_id in TASK_ORDER},
             difficulty_scores={"easy": 0.0, "medium": 0.0, "hard": 0.0},
@@ -410,7 +279,8 @@ class ITMentalHealthEnvironment:
         task = TASK_ORDER[0]
         scenario_text, _ = self._current_scenarios[task]
         return MentalHealthObservation(
-            scenario=scenario_text, feedback=f"New episode (seed={actual_seed}).",
+            scenario=scenario_text,
+            feedback=f"New episode (seed={actual_seed}).",
             task_id=task,
             metadata={
                 "seed": actual_seed,
@@ -420,6 +290,7 @@ class ITMentalHealthEnvironment:
                 "task_cumulative_score": 0.0,
                 "difficulty_cumulative_score": 0.0,
                 "overall_cumulative_score": 0.0,
+                "task_id": task,
             },
         )
 
@@ -427,7 +298,6 @@ class ITMentalHealthEnvironment:
         if not self._state.episode_id or not self._current_scenarios:
             raise ValueError("Episode not initialized. Call /reset before /step.")
 
-        # Guard: refuse to grade once the episode is finished.
         if self._task_index >= len(TASK_ORDER):
             observation = MentalHealthObservation(
                 scenario="Episode already finished. Call /reset to start a new one.",
@@ -441,12 +311,10 @@ class ITMentalHealthEnvironment:
         task = self._state.current_task
         if action.task_id != task:
             raise ValueError(f"task_id mismatch: expected '{task}', got '{action.task_id}'.")
-        scenario_text, ground_truth = self._current_scenarios.get(task, ("", {}))
 
-        reward, breakdown, feedback = _llm_judge_grade(
-            task_id=task, scenario=scenario_text,
-            response=action.response, ground_truth=ground_truth,
-        )
+        _, ground_truth = self._current_scenarios.get(task, ("", {}))
+        reward, breakdown, feedback = grade_response(task, action.response, ground_truth)
+
         score_metadata = self._score_metadata(task, reward)
         self._state.cumulative_reward += reward
         self._state.task_scores[task] = score_metadata["task_cumulative_score"]

openenv.yaml

@@ -6,8 +6,8 @@ description: >
   A reinforcement learning environment for mental health assessment and
   intervention planning in IT and software engineering teams. Agents solve
   three progressively harder tasks: burnout detection, stress triage, and
-  team-level intervention planning. Each task returns a normalized reward in
-  the 0.0-1.0 range based on rubric coverage, structure, and clinical accuracy.
+  team-level intervention planning. Each task has a deterministic programmatic
+  grader and returns a normalized reward in the 0.0-1.0 range.
 
 type: typed
 
@@ -25,6 +25,7 @@ tasks:
   - id: burnout_detection
     difficulty: easy
     description: Identify Maslach burnout dimensions and severity from an employee profile
+    objective: Identify burnout dimensions, severity, top red flags, and HR escalation need from a single employee profile
     reward_range: [0.0, 1.0]
     grader: burnout_detection_grader
     graders:
@@ -33,6 +34,7 @@ tasks:
   - id: stress_triage
     difficulty: medium
     description: Triage 3 IT employees by stress tier and recommend immediate actions
+    objective: Classify three employee cases by urgency, rank intervention priority, and recommend immediate plus 2-week support
     reward_range: [0.0, 1.0]
     grader: stress_triage_grader
     graders:
@@ -41,6 +43,7 @@ tasks:
   - id: intervention_plan
     difficulty: hard
     description: Design a 4-week mental health intervention plan for a burned-out IT team
+    objective: Create a four-week intervention plan with owners, KPIs, risk, and budget for a burned-out software team
     reward_range: [0.0, 1.0]
     grader: intervention_plan_grader
     graders:

server/app.py

@@ -66,16 +66,19 @@ TASK_METADATA = {
     "burnout_detection": {
         "difficulty": "easy",
         "description": "Identify Maslach burnout dimensions, severity, red flags, and whether HR escalation is needed.",
+        "objective": "Assess one employee profile and determine burnout dimensions, severity, red flags, and escalation need.",
         "graders": ["burnout_detection_grader"],
     },
     "stress_triage": {
         "difficulty": "medium",
         "description": "Triage three IT employees by stress tier and recommend immediate and medium-term support.",
+        "objective": "Classify three employees by urgency, rank intervention priority, and recommend immediate plus 2-week support.",
         "graders": ["stress_triage_grader"],
     },
     "intervention_plan": {
         "difficulty": "hard",
         "description": "Design a four-week intervention plan for a software team facing systemic burnout.",
+        "objective": "Produce a four-week intervention plan with owners, measurable outcomes, KPIs, risk, and budget.",
         "graders": ["intervention_plan_grader"],
     },
 }
@@ -258,6 +261,7 @@ def list_tasks():
                 "id": tid,
                 "difficulty": TASK_METADATA[tid]["difficulty"],
                 "description": TASK_METADATA[tid]["description"],
+                "objective": TASK_METADATA[tid]["objective"],
                 "grader": TASK_METADATA[tid]["graders"][0],
                 "graders": TASK_METADATA[tid]["graders"],
                 "scoring": {

tasks.py

@@ -5,23 +5,27 @@ Static task registry for OpenEnv task discoverability.
 TASKS = [
     {
         "id": "burnout_detection",
+        "difficulty": "easy",
+        "objective": "Identify burnout dimensions, severity, top red flags, and whether immediate HR escalation is needed.",
         "grader": "burnout_detection_grader",
         "graders": ["burnout_detection_grader"],
     },
     {
         "id": "stress_triage",
+        "difficulty": "medium",
+        "objective": "Classify three employee cases by stress tier, rank urgency, and recommend immediate plus 2-week support.",
         "grader": "stress_triage_grader",
         "graders": ["stress_triage_grader"],
     },
     {
         "id": "intervention_plan",
+        "difficulty": "hard",
+        "objective": "Create a four-week intervention plan with owners, KPIs, risk, and budget for a burned-out team.",
         "grader": "intervention_plan_grader",
         "graders": ["intervention_plan_grader"],
     },
 ]
 
-
 TASK_IDS = [task["id"] for task in TASKS]
 
-
 __all__ = ["TASKS", "TASK_IDS"]