Fixing state leak bug

backend/tracefix_rl_environment.py

@@ -20,7 +20,7 @@ class TraceFixRLEnvironment(Environment):
         self._gym = TraceFixRLGym()
         self._state = State(episode_id="", step_count=0)
 
-    def reset(self, difficulty: str | None = None) -> CodeObservation:
+    def reset(self, difficulty: str | None = None, task_name: str | None = None) -> CodeObservation:
         if difficulty == "easy":
             self._gym.training_step = 1
         elif difficulty == "medium":
@@ -28,7 +28,18 @@ class TraceFixRLEnvironment(Environment):
         elif difficulty == "hard":
             self._gym.training_step = 6000
 
-        obs, system_prompt = self._gym.reset()
+        task_dict = None
+        if task_name and task_name != "tracefix_rl":
+            try:
+                from tasks.tasks import ALL_TASKS
+                for t in ALL_TASKS:
+                    if t.get("name") == task_name:
+                        task_dict = t
+                        break
+            except ImportError:
+                pass
+
+        obs, system_prompt = self._gym.reset(task_index=task_dict)
         self._state = State(
             episode_id=obs.info.get("episode_id", ""),
             step_count=obs.step_count,

inference.py

@@ -325,11 +325,13 @@ async def run(difficulty: Optional[str] = None, show_thought: bool = False) -> N
         else:
             env = TraceFixRLEnv(base_url=ENV_BASE_URL)
 
+        reset_kwargs = {}
         if difficulty:
-            reset_kwargs = {"difficulty": difficulty}
-            result = await env.reset(**reset_kwargs)
-        else:
-            result = await env.reset()
+            reset_kwargs["difficulty"] = difficulty
+        if TASK_NAME and TASK_NAME != "tracefix_rl":
+            reset_kwargs["task_name"] = TASK_NAME
+
+        result = await env.reset(**reset_kwargs)
         task_name = result.observation.info.get("task_name") or TASK_NAME
         log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
         started = True

vision_ui.py

@@ -287,6 +287,16 @@ def _build_env(
 
 
 def sync_tasks(selected, grid_name):
+    if not selected:
+        return (
+            gr.skip(), 
+            gr.skip(), 
+            gr.skip(), 
+            gr.skip(), 
+            gr.skip(), 
+            gr.skip(),
+            gr.skip()
+        )
     if grid_name == "easy":
         easy_val = selected
         med_val = None