Upload server/overview_environment.py with huggingface_hub

server/overview_environment.py

@@ -0,0 +1,99 @@
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import Action, Observation, State
+
+import random
+from typing import Any, Dict, Optional
+from dataclasses import dataclass, field
+from uuid import uuid4
+
+from overview_env.models import OverviewObservation, OverviewAction, AnalysisResult, TaskType
+from overview_env.tasks.definitions import (
+    get_all_tasks,
+    get_task_by_id,
+    OverviewTask,
+    OverviewTaskEvaluator,
+    GradingResult,
+)
+
+
+@dataclass
+class _EpisodeState:
+    task: OverviewTask
+    episode_id: str
+    current_step: int = 0
+    cumulative_reward: float = 0.0
+    submitted_analysis: Dict[str, Any] = field(default_factory=dict)
+    episode_complete: bool = False
+
+
+DIFFICULTY_WEIGHTS = {"easy": 0.15, "medium": 0.12, "hard": 0.08}
+
+
+class OverviewEnvironment(Environment):
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self, task_id: Optional[str] = None, seed: Optional[int] = None, max_steps: int = 10):
+        self._task_id = task_id
+        self._seed = seed
+        self._max_steps = max_steps
+        self._ep: Optional[_EpisodeState] = None
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        if seed is not None:
+            random.seed(seed)
+
+    def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, task_id: Optional[str] = None, **kwargs: Any) -> Observation:
+        if seed is not None:
+            random.seed(seed)
+        target_task_id = task_id or self._task_id
+        if target_task_id:
+            task = get_task_by_id(target_task_id)
+        else:
+            task = random.choice(get_all_tasks())
+        eid = episode_id or str(uuid4())
+        self._ep = _EpisodeState(task=task, episode_id=eid, current_step=0, cumulative_reward=0.0, submitted_analysis={}, episode_complete=False)
+        self._state = State(episode_id=eid, step_count=0)
+        obs = self._build_observation()
+        return Observation(done=False, reward=0.0, metadata=obs.model_dump())
+
+    def step(self, action: Action, timeout_s: Optional[float] = None, **kwargs: Any) -> Observation:
+        if self._ep is None:
+            return Observation(done=True, reward=0.0, metadata={"error": "Environment not reset. Call reset() first."})
+        if self._ep.episode_complete:
+            return Observation(done=True, reward=0.0, metadata={"error": "Episode already finished."})
+        action_data = {}
+        if hasattr(action, "data") and isinstance(action.data, dict):
+            action_data = action.data
+        elif isinstance(action, dict):
+            action_data = action
+        elif hasattr(action, "__dict__"):
+            action_data = vars(action)
+        try:
+            env_action = OverviewAction.model_validate(action_data)
+        except Exception:
+            env_action = OverviewAction(analysis=AnalysisResult())
+        self._ep.current_step += 1
+        self._state.step_count = self._ep.current_step
+        analysis_dict = env_action.analysis.model_dump()
+        self._ep.submitted_analysis = analysis_dict
+        evaluator = OverviewTaskEvaluator(self._ep.task)
+        result = evaluator.grade(analysis_dict)
+        reward_value = result.score * DIFFICULTY_WEIGHTS.get(self._ep.task.difficulty, 0.1)
+        terminal = env_action.submit or self._ep.current_step >= self._ep.task.max_steps
+        if terminal:
+            final_score = result.score
+            if self._ep.current_step <= self._ep.task.max_steps * 0.5:
+                final_score += 0.1
+            reward_value = min(1.0, final_score)
+            self._ep.episode_complete = True
+        self._ep.cumulative_reward += reward_value
+        obs = self._build_observation()
+        step_result = {"observation": obs.model_dump(), "reward": {"value": reward_value, "total": result.score, "feedback": result.feedback}, "done": self._ep.episode_complete, "info": {"step": self._ep.current_step}}
+        return Observation(done=self._ep.episode_complete, reward=reward_value, metadata=step_result)
+
+    @property
+    def state(self) -> State:
+        return self._state
+
+    def _build_observation(self) -> OverviewObservation:
+        ep = self._ep
+        return OverviewObservation(task_id=ep.task.task_id, task_type=TaskType(ep.task.task_type), task_name=ep.task.name, task_description=ep.task.description, difficulty=ep.task.difficulty, input_text=ep.task.input_text, question=ep.task.question, expected_output=ep.task.expected_output, max_steps=ep.task.max_steps)