Deva22467/govt-scheme-openenv

1import random2import numpy as np3from typing import Dict, Any, Tuple, List4from data import PROFILES, SCHEMES_DB, get_profiles_by_task5from models import Observation, Action6 7class SchemeEnv:8    """9    OpenEnv Reinforcement Learning Multi-Step trajectory environment for Government Scheme Recommendation.10    Models decision-making used by digital governance assistants to improve scheme adoption.11    """12    def __init__(self, seed: int = 42):13        self.profiles = PROFILES14        self.current_profile: Dict[str, Any] = {}15        self.task_mode = "medium"16        self.stage = 017        self.episode_reward = 0.018        self.seed(seed)19 20    def seed(self, seed: int):21        """Sets the seed for reproducible deterministic states."""22        random.seed(seed)23        np.random.seed(seed)24 25    def reset(self, task: str = None, profile_id: str = None) -> Observation:26        """Resets the environment."""27        self.stage = 028        self.episode_reward = 0.029        30        if profile_id:31            candidates = [p for p in self.profiles if p['id'] == profile_id]32            self.current_profile = candidates[0] if candidates else random.choice(self.profiles)33        elif task:34            candidates = get_profiles_by_task(task)35            self.current_profile = random.choice(candidates) if candidates else random.choice(self.profiles)36        else:37            self.current_profile = random.choice(self.profiles)38            39        self.task_mode = self.current_profile['task']40        return self.state()41 42    def state(self) -> Observation:43        """Returns Observation class including dynamic noise."""44        if not self.current_profile:45            # Fallback initialization bounding46            return Observation(age=0, income=0, category="", occupation="", state="", task="", stage=0)47        48        state_dict = self.current_profile['profile'].copy()49        50        if self.task_mode == "hard":51            if "occupation" in state_dict:52                noise_postfix = random.choice([" + part time freelancer", " / informal gig worker", " + hobbyist"])53                state_dict["occupation"] = state_dict["occupation"] + noise_postfix54                55            irrelevant_noise = [56                {"has_pets": random.choice([True, False])},57                {"favorite_color": random.choice(["Blue", "Green", "Red"])},58                {"zodiac_sign": random.choice(["Libra", "Aries", "Taurus", "Gemini"])}59            ]60            state_dict.update(random.choice(irrelevant_noise))61 62        state_dict["task"] = self.task_mode63        state_dict["stage"] = self.stage64        return Observation(**state_dict)65 66    def step(self, action: Action) -> Tuple[Observation, float, bool, Dict]:67        """Trajectory Step. Receives structured Action Schema Pydantic."""68        if not self.current_profile:69            raise ValueError("Environment must be reset.")70 71        gt_schemes = set(self.current_profile['eligible_schemes'])72        gt_best = self.current_profile['best_scheme']73        gt_reason = self.current_profile['reasoning']74        75        step_reward = 0.076        info = {77            "stage": self.stage,78            "task": self.task_mode,79            "eligible_score": 0.0,80            "best_score": 0.0,81            "reason_score": 0.0,82            "total_reward": 0.083        }84        85        if self.stage == 0:86            pred_schemes = set(action.schemes) if action.schemes else set()87            correct_schemes = gt_schemes.intersection(pred_schemes)88            89            if len(gt_schemes) > 0:90                fraction = len(correct_schemes) / len(gt_schemes)91                step_reward += fraction * 0.392            info["eligible_score"] = step_reward93            94        elif self.stage == 1:95            pred_best = action.best if action.best else ''96            if pred_best == gt_best:97                step_reward += 0.498                99            if pred_best in gt_schemes:100                benefit_values = [SCHEMES_DB.get(s, {}).get("benefit_amount", 0) for s in gt_schemes]101                if benefit_values:102                    max_benefit = max(benefit_values)103                    pred_benefit = SCHEMES_DB.get(pred_best, {}).get("benefit_amount", 0)104                    if max_benefit > 0 and pred_benefit == max_benefit:105                        step_reward += 0.2106            info["best_score"] = step_reward107            108        elif self.stage == 2:109            pred_reason = action.reasoning if action.reasoning else ''110            if pred_reason and isinstance(pred_reason, str) and len(pred_reason.split()) >= 3:111                gt_words = set(gt_reason.lower().split())112                pred_words = set(pred_reason.lower().split())113                if len(gt_words.intersection(pred_words)) >= 2:114                    step_reward += 0.3115            info["reason_score"] = step_reward116        117        allowed_reward = 1.0 - self.episode_reward118        yield_reward = min(step_reward, allowed_reward)119        self.episode_reward += yield_reward120        self.stage += 1121        done = (self.stage >= 3)122        info["total_reward"] = self.episode_reward123        124        yield_reward = max(0.05, min(yield_reward, 0.95))125        126        return self.state(), yield_reward, done, info127