""" DriftEnv — app.py """ import json import os import random from typing import Optional from pydantic import BaseModel from openenv.core import Action, Observation, State, Environment DATASET_PATH = os.path.join(os.path.dirname(__file__), "dataset.json") with open(DATASET_PATH, "r") as f: SCENARIOS = json.load(f) class DriftEnvObservation(BaseModel): instruction: str context_shift: Optional[str] = None step: int history: list done: bool task: str class DriftEnvAction(BaseModel): response: str class DriftEnvReward(BaseModel): reward: float final_score: Optional[float] = None phase: str feedback: str _state = { "scenario": None, "task": "medium", "step_count": 0, "shift_triggered": False, "interpretation_score": None, "pivot_score": None, "completion_score": None, "done": False, "history": [], "prev_responses": [], # agent responses per step, used for pivot/no_stale scoring } # --------------------------------------------------------------------------- # Reward component helpers # --------------------------------------------------------------------------- def _score_format(response: str) -> float: """R_format: reward concise, structured responses.""" n = len(response) if n < 200: return 1.0 if n < 500: return 0.5 return 0.0 def _extract_unique_keywords(target_text: str, exclusion_text: str) -> list: """Return words >4 letters in target_text that don't appear in exclusion_text. Prevents agents from gaming the score by echoing words from the visible instruction.""" excl = set(exclusion_text.lower().split()) return [w for w in target_text.lower().split() if len(w) > 4 and w not in excl] def _score_interpretation(response: str, hidden_interpretation: str, exclusion_text: str = "") -> float: """R_interpretation: keyword overlap with hidden_interpretation, excluding words already visible in the instruction/context_shift.""" resp_lower = response.lower() keywords = _extract_unique_keywords(hidden_interpretation, exclusion_text) if not keywords: return 0.0 hits = sum(1 for kw in keywords if kw in resp_lower) # Linear scale: ratio >= 0.4 → 1.0, proportional below return round(min(hits / len(keywords) / 0.4, 1.0), 4) def _score_pivot(response: str, correct_pivot: str, step1_response: Optional[str], exclusion_text: str = "") -> float: """R_pivot: on step >= 2, keyword overlap with correct_pivot (excluding visible words) AND lexical distance from the agent's own step-1 response.""" if step1_response is None: return 0.0 resp_lower = response.lower() # (a) keyword overlap with correct pivot — unique words only keywords = _extract_unique_keywords(correct_pivot, exclusion_text) if keywords: hits = sum(1 for kw in keywords if kw in resp_lower) kw_score = min(hits / len(keywords) / 0.4, 1.0) else: kw_score = 0.0 # (b) lexical distance from step-1 response (0 = identical, 1 = fully different) prev_words = set(w for w in step1_response.lower().split() if len(w) > 3) curr_words = set(w for w in resp_lower.split() if len(w) > 3) if prev_words: shared = len(prev_words & curr_words) lexical_dist = 1.0 - (shared / len(prev_words)) else: lexical_dist = 1.0 return round((kw_score + lexical_dist) / 2, 4) def _score_no_stale(response: str, wrong_pivots: list, step1_response: Optional[str]) -> float: """R_no_stale: on step >= 2, penalise responses that match wrong_pivots or are too similar to the agent's step-1 response (anti-reward-hacking signal).""" if step1_response is None: return 1.0 # step 1 — nothing can be stale yet resp_lower = response.lower() # Hard zero if response echoes a known wrong pivot wrong_match = any(w.lower()[:30] in resp_lower for w in wrong_pivots) if wrong_match: return 0.0 # Graded penalty for repeating step-1 content prev_words = set(w for w in step1_response.lower().split() if len(w) > 3) curr_words = set(w for w in resp_lower.split() if len(w) > 3) if prev_words: shared = len(prev_words & curr_words) similarity = shared / len(prev_words) if similarity > 0.7: return 0.0 if similarity > 0.4: return round(1.0 - similarity, 4) return 1.0 def _compute_reward( action: str, scenario: dict, step1_response: Optional[str], ) -> tuple[float, dict]: """Compute all 4 reward components and return (weighted_total, components_dict).""" # Words visible to the agent — excluded from keyword pools to close the echo exploit exclusion_text = scenario["initial_instruction"] + " " + (scenario.get("context_shift") or "") r_fmt = _score_format(action) r_interp = _score_interpretation(action, scenario["hidden_interpretation"], exclusion_text) r_pivot = _score_pivot(action, scenario["correct_pivot"], step1_response, exclusion_text) r_stale = _score_no_stale(action, scenario.get("wrong_pivots", []), step1_response) total = round(0.1 * r_fmt + 0.3 * r_interp + 0.4 * r_pivot + 0.2 * r_stale, 4) components = { "R_format": round(r_fmt, 4), "R_interpretation": round(r_interp, 4), "R_pivot": round(r_pivot, 4), "R_no_stale": round(r_stale, 4), } return total, components # --------------------------------------------------------------------------- # Core env logic # --------------------------------------------------------------------------- def _observe() -> dict: s = _state["scenario"] return { "instruction": s["initial_instruction"], "context_shift": s["context_shift"] if _state["shift_triggered"] else None, "step": _state["step_count"], "history": _state["history"], "done": _state["done"], "task": _state["task"], } def reset(task: str = "medium", holdout_only: bool = False) -> dict: global _state pool = [s for s in SCENARIOS if s["holdout"]] if holdout_only else [s for s in SCENARIOS if not s["holdout"]] scenario = random.choice(pool) _state = { "scenario": scenario, "task": task, "step_count": 0, "shift_triggered": False, "interpretation_score": None, "pivot_score": None, "completion_score": None, "done": False, "history": [], "prev_responses": [], } return _observe() def step(action: str) -> dict: if _state["done"]: return { "observation": _observe(), "reward": 0.0, "final_score": None, "done": True, "info": {"error": "Episode finished. Call reset() to start a new episode."}, } s = _state["scenario"] task = _state["task"] _state["step_count"] += 1 n = _state["step_count"] # step-1 response is the pivot/no_stale reference for all subsequent steps step1_response = _state["prev_responses"][0] if _state["prev_responses"] else None _state["prev_responses"].append(action) reward, components = _compute_reward(action, s, step1_response) final_score = None phase = "" if n == 1: phase = "interpretation" _state["interpretation_score"] = reward if task == "easy": _state["done"] = True final_score = reward else: _state["shift_triggered"] = True elif n == 2: phase = "pivot" _state["pivot_score"] = reward if task == "medium": scores = [_state["interpretation_score"], _state["pivot_score"]] final_score = round(sum(scores) / len(scores), 4) _state["done"] = True elif n == 3: phase = "completion" # Hard mode: agent must score well on BOTH R_interpretation (remembered # original intent) and R_pivot (executed the shift) simultaneously. # The weighted formula already enforces this — no separate code needed. _state["completion_score"] = reward scores = [_state["interpretation_score"], _state["pivot_score"], _state["completion_score"]] final_score = round(sum(scores) / len(scores), 4) _state["done"] = True _state["history"].append({ "step": n, "phase": phase, "action": action[:300], "reward": reward, "rewards": components, }) info = { "phase": phase, "task": task, "rewards": components, } if _state["done"] and final_score is not None: info["final_score"] = final_score info["breakdown"] = { "interpretation": _state["interpretation_score"], "pivot": _state["pivot_score"], "completion": _state["completion_score"], } return { "observation": _observe(), "reward": reward, "final_score": final_score, "done": _state["done"], "info": info, } def state() -> dict: s = _state["scenario"] return { "scenario_id": s["id"] if s else None, "domain": s["domain"] if s else None, "task": _state["task"], "step_count": _state["step_count"], "shift_triggered": _state["shift_triggered"], "scores": { "interpretation": _state["interpretation_score"], "pivot": _state["pivot_score"], "completion": _state["completion_score"], }, "done": _state["done"], "history": _state["history"], } # --- OpenEnv Environment subclass --- class _DriftEnvAction(Action): response: str class _DriftEnvObservation(Observation): instruction: str context_shift: Optional[str] = None step: int = 0 history: list = [] task: str = "medium" class _DriftEnvState(State): scenario_id: Optional[int] = None domain: Optional[str] = None task: str = "medium" shift_triggered: bool = False scores: dict = {} history: list = [] class DriftEnvironment(Environment): def __init__(self): super().__init__() def _observe(self) -> _DriftEnvObservation: obs = _observe() return _DriftEnvObservation( instruction=obs["instruction"], context_shift=obs["context_shift"], step=obs["step"], history=obs["history"], done=obs["done"], task=obs["task"], ) def reset(self, seed=None, episode_id=None, **kwargs) -> _DriftEnvObservation: task = kwargs.get("task", "medium") holdout_only = bool(kwargs.get("holdout_only", False)) reset(task=task, holdout_only=holdout_only) return self._observe() def step(self, action: _DriftEnvAction, timeout_s=None, **kwargs) -> _DriftEnvObservation: result = step(action.response) obs = self._observe() final_score = result.get("final_score") obs.reward = final_score if final_score is not None else result.get("reward", 0.0) return obs @property def state(self) -> _DriftEnvState: s = state() return _DriftEnvState( scenario_id=s["scenario_id"], domain=s["domain"], task=s["task"], step_count=s["step_count"], shift_triggered=s["shift_triggered"], scores=s["scores"], history=s["history"], ) def create_health_app(app): @app.get("/") async def root(): return { "name": "DriftEnv", "status": "running", "description": "RL environment for testing AI agent robustness under ambiguity and context shift", "tasks": ["easy", "medium", "hard"], "endpoints": { "reset": "POST /reset", "step": "POST /step", "state": "GET /state", "docs": "GET /docs" } } return app def main(): import uvicorn from openenv.core import create_app app = create_app( env=DriftEnvironment, action_cls=_DriftEnvAction, observation_cls=_DriftEnvObservation, env_name="driftenv", ) create_health_app(app) # register routes uvicorn.run(app, host="0.0.0.0", port=7860) if __name__ == "__main__": main()