match reference env pattern exactly

sql_env/server.py

@@ -18,9 +18,9 @@ except ImportError:
 
 
 class SQLCorrectionEnvironment(Environment):
+    SUPPORTS_CONCURRENT_SESSIONS = True
 
     def __init__(self):
-        super().__init__()
         self._difficulty = "easy"
         self._current_task = None
         self._step_count = 0
@@ -28,8 +28,12 @@ class SQLCorrectionEnvironment(Environment):
         self._last_reward = 0.0
         self._rewards_history = []
 
-    def reset(self, difficulty: str = "easy", task_id: str = None, **kwargs) -> SQLObservation:
-        actual_difficulty = task_id or difficulty or "easy"
+    def reset(self, seed=None, episode_id=None, **kwargs) -> SQLObservation:
+        actual_difficulty = (
+            kwargs.get("task_id") or
+            kwargs.get("difficulty") or
+            "easy"
+        )
         self._difficulty = actual_difficulty
         tasks = TASK_SETS.get(actual_difficulty, TASK_SETS["easy"])
         self._current_task = random.choice(tasks)
@@ -72,6 +76,50 @@ class SQLCorrectionEnvironment(Environment):
             done=done,
         )
 
+    @property
+    def state(self) -> SQLState:
+        if self._current_task is None:
+            return SQLState(
+                task_id="none",
+                difficulty="none",
+                step_count=0,
+                max_steps=0,
+                done=False,
+                last_reward=0.0,
+                rewards_history=[],
+            )
+        return SQLState(
+            task_id=self._current_task.task_id,
+            difficulty=self._difficulty,
+            step_count=self._step_count,
+            max_steps=self._current_task.max_steps,
+            done=self._done,
+            last_reward=self._last_reward,
+            rewards_history=self._rewards_history,
+        )
+    def step(self, action: SQLAction) -> SQLObservation:
+        if self._current_task is None:
+            self.reset()
+        self._step_count += 1
+        reward_obj = grade(action, self._current_task)
+        reward = reward_obj.value
+        self._last_reward = reward
+        self._rewards_history.append(reward)
+        done = (reward >= 0.95) or (self._step_count >= self._current_task.max_steps)
+        self._done = done
+        feedback = generate_feedback(action, self._current_task, reward_obj)
+        return SQLObservation(
+            task_id=self._current_task.task_id,
+            broken_query=self._current_task.broken_query,
+            schema_context=self._current_task.schema_context,
+            error_hint=self._current_task.error_hint,
+            step_number=self._step_count,
+            previous_attempt=action.corrected_query,
+            feedback=feedback,
+            reward=reward,
+            done=done,
+        )
+
     def state(self) -> SQLState:
         if self._current_task is None:
             return SQLState(
@@ -145,7 +193,6 @@ app = create_app(
     SQLAction,
     SQLObservation,
     env_name="sql-correction-env",
-    max_concurrent_envs=1,
 )