good to go

README.md

@@ -11,29 +11,103 @@ tags:
 ---
 # OpenEnv SQL Data Analyst
 
-## The Goal
-The objective is straightforward: the AI must interpret an internal database schema and translate human business instructions into valid SQL queries natively executed against the pipeline. 
-
-Tasks increase in difficulty, pushing models to understand complex relationships such as:
-- Simple Table Retrievals (`SELECT`, `WHERE`)
-- Aggregating related tables (`JOIN`, `GROUP BY`)
-- Complex constraint execution bounds (Logic matching)
-
-## Fast Execution
-1. **Install:** `pip install -r requirements.txt`
-2. **Setup your preferred engine:**
-   ```powershell
-   $env:GEMINI_API_KEY="your-api-key" 
-   # Or use OPENAI_API_KEY 
+## Environment Description and Motivation
+This environment simulates a real business workflow: a data analyst converting stakeholder requests into executable SQL over an internal company warehouse.
+
+Why this is real-world:
+- HR and finance teams routinely need ad-hoc SQL analysis.
+- Models must handle schema understanding, joins, aggregation, and constraints.
+- Evaluation is done against deterministic graders over actual query results (not string matching).
+
+The environment implements the OpenEnv lifecycle (`reset()`, `step()`, `state()`) and is deployable on Hugging Face Spaces as a containerized service.
+
+## Action Space
+`SqlAction`
+- `query: str` - a SQL query to execute against the provided SQLite schema.
+
+## Observation Space
+`SqlObservation`
+- `current_task_instruction: str` - current objective text.
+- `schema_info: str` - schema description.
+- `task_id: str` - stable task identifier.
+- `difficulty: str` - `easy | medium | hard`.
+- `execution_result: Optional[str]` - previous query result rows as JSON.
+- `execution_error: Optional[str]` - previous execution or safety error.
+- `task_score: float` - grader output in `[0.0, 1.0]`.
+- `grader_feedback: Optional[str]` - deterministic grader feedback.
+
+## State
+`SqlState`
+- Tracks episode id, step count, current task index, accumulated reward, per-task scores, and attempts.
+- Exposed via `state()` for inspection and reproducibility.
+
+## Tasks and Difficulty Progression
+1. **Easy - `employee_payroll_overview`**  
+   List employee names and salaries in descending salary order.
+2. **Medium - `department_budget_summary`**  
+   Compute average salary per department with required output schema.
+3. **Hard - `senior_engineering_comp_review`**  
+   Multi-table logic to identify engineering employees meeting senior salary thresholds.
+
+All tasks have deterministic graders producing scores from `0.0` to `1.0`.
+
+## Reward Function
+Reward is shaped for trajectory-level signal (not binary terminal only):
+- `+0.1` valid SQL execution bonus
+- `+0.9 * task_score` progress toward correctness
+- `-0.02` per step (discourages loops)
+- `-0.5` invalid query penalty
+- `-1.0` safety penalty for destructive SQL (`DROP/DELETE/TRUNCATE/ALTER/UPDATE/INSERT`)
+
+Task completion threshold is `task_score >= 0.95`.
+Episode ends when all tasks are solved or max attempts for a task are exhausted.
+
+## Setup and Usage
+1. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+2. Set API key for baseline:
+   ```bash
+   export OPENAI_API_KEY="your-key"
    ```
-3. **Run the Simulation Base:**
-   ```powershell
+3. Run baseline:
+   ```bash
    python inference.py
    ```
 
-## Architecture Workflow
+The baseline uses the OpenAI API client with deterministic settings:
+- `temperature=0.0`
+- fixed seed (`OPENAI_SEED`, default `42`)
+- fixed task order and deterministic grader
+
+## Baseline Score Reporting
+`inference.py` prints:
+- per-task best score
+- average task score
+- episode return
+- model and seed used
+
+This creates reproducible benchmark records for easy/medium/hard tasks.
+
+## Docker and Hugging Face Space
+- `Dockerfile` is included and starts `uvicorn server.app:app`.
+- Space metadata is configured in this `README.md` frontmatter with `sdk: docker`.
+- Health endpoint: `/health`
 
-This environment follows strict [OpenEnv Gym interface](https://github.com/meta-pytorch/OpenEnv) standards, utilizing `step()`, `reset()`, and `state()` endpoints wrapped locally or exposed natively via Hugging Face.
+Local container run:
+```bash
+docker build -t sql-agent-env .
+docker run -p 7860:7860 sql-agent-env
+```
+
+## OpenEnv Metadata
+OpenEnv runtime metadata is declared in `openenv.yaml`.
+
+Validation command:
+```bash
+openenv validate
+```
 
 ```mermaid
 flowchart TD
@@ -69,4 +143,4 @@ flowchart TD
 To read more about exactly how the logic is handled under the hood, access the OOP components inside the root directory:
 - `server/environment.py`: The orchestrator handling the OpenEnv lifecycle.
 - `database/sqlite_manager.py`: Creates ephemeral SQLite sandboxes per session.
-- `core/grader.py`: Evaluates logic correctness independent of string semantics by leveraging Pandas DataFrame alignments.
+- `core/grader.py`: Deterministic result grader with partial credit and detailed feedback.

core/grader.py

@@ -1,4 +1,6 @@
 import pandas as pd
+from collections import Counter
+from typing import Dict, Any
 
 class DataFrameGrader:
     """Handles logic for comparing agent SQL result DataFrames against expected DataFrames."""
@@ -10,29 +12,62 @@ class DataFrameGrader:
         Checks for exact match (columns, rows, data types).
         Falls back to unordered match if the structure is the same but sorting is off.
         """
+        return DataFrameGrader.grade_with_details(agent_df, expected_df)["score"]
+
+    @staticmethod
+    def grade_with_details(agent_df: pd.DataFrame, expected_df: pd.DataFrame) -> Dict[str, Any]:
+        """
+        Deterministic grader with partial signals:
+        - 0.4 for column correctness
+        - 0.2 for row count closeness
+        - 0.4 for row content match (multiset, order-insensitive)
+        """
         try:
             agent_norm = agent_df.reset_index(drop=True).sort_index(axis=1)
             expected_norm = expected_df.reset_index(drop=True).sort_index(axis=1)
-            
-            # Fast exact comparison
-            if set(agent_norm.columns) == set(expected_norm.columns):
-                
-                # Check directly first
-                if agent_norm.equals(expected_norm):
-                    return 1.0
-                
-                # Re-align column order if mismatched
-                agent_aligned = agent_norm[expected_norm.columns]
-                if agent_aligned.equals(expected_norm):
-                    return 1.0
-                
-                # Check unordered row set equivalence
-                agent_tuples = set(tuple(row) for row in agent_aligned.to_numpy())
-                expected_tuples = set(tuple(row) for row in expected_norm.to_numpy())
-                
-                if len(agent_tuples) == len(expected_tuples) and agent_tuples == expected_tuples:
-                    return 0.8  # Unordered match yields partial points
-            
-            return 0.0
-        except Exception:
-            return 0.0
+
+            expected_cols = list(expected_norm.columns)
+            agent_cols = list(agent_norm.columns)
+            same_col_set = set(agent_cols) == set(expected_cols)
+            col_score = 0.4 if same_col_set else 0.0
+
+            if not same_col_set:
+                return {
+                    "score": 0.0,
+                    "feedback": f"Column mismatch. Expected {expected_cols}, got {agent_cols}.",
+                    "details": {"column_score": 0.0, "row_count_score": 0.0, "row_content_score": 0.0},
+                }
+
+            agent_aligned = agent_norm[expected_cols]
+
+            expected_count = len(expected_norm)
+            agent_count = len(agent_aligned)
+            if expected_count == 0:
+                row_count_score = 0.2 if agent_count == 0 else 0.0
+            else:
+                row_delta = abs(agent_count - expected_count)
+                row_count_score = 0.2 * max(0.0, 1.0 - (row_delta / expected_count))
+
+            agent_rows = Counter(tuple(row) for row in agent_aligned.to_numpy())
+            expected_rows = Counter(tuple(row) for row in expected_norm.to_numpy())
+            overlap = sum((agent_rows & expected_rows).values())
+            denom = max(1, sum(expected_rows.values()))
+            row_content_score = 0.4 * (overlap / denom)
+
+            score = round(min(1.0, col_score + row_count_score + row_content_score), 4)
+            feedback = "Exact match." if score == 1.0 else "Partially correct result."
+            return {
+                "score": score,
+                "feedback": feedback,
+                "details": {
+                    "column_score": round(col_score, 4),
+                    "row_count_score": round(row_count_score, 4),
+                    "row_content_score": round(row_content_score, 4),
+                },
+            }
+        except Exception as exc:
+            return {
+                "score": 0.0,
+                "feedback": f"Grader failed: {exc}",
+                "details": {"column_score": 0.0, "row_count_score": 0.0, "row_content_score": 0.0},
+            }

core/tasks.py

@@ -3,7 +3,9 @@ from typing import List
 
 @dataclass
 class SqlTask:
+    task_id: str
     difficulty: int
+    difficulty_label: str
     instruction: str
     expected_query: str
 
@@ -12,13 +14,17 @@ class EnvironmentTasks:
     def __init__(self):
         self.tasks: List[SqlTask] = [
             SqlTask(
+                task_id="employee_payroll_overview",
                 difficulty=1,
-                instruction="EASY: Find all names and salaries of employees, ordering by their salaries descending.",
+                difficulty_label="easy",
+                instruction="EASY: As an HR analyst, list all employee names and salaries ordered by salary descending for payroll review.",
                 expected_query="SELECT NAME, SALARY FROM EMPLOYEES ORDER BY SALARY DESC"
             ),
             SqlTask(
+                task_id="department_budget_summary",
                 difficulty=2,
-                instruction="MEDIUM: Return the department names alongside the average salary per department. Columns must be 'DEPARTMENT_NAME' and 'AVG_SALARY'.",
+                difficulty_label="medium",
+                instruction="MEDIUM: As a finance analyst, return each department name with average salary. Output columns must be 'DEPARTMENT_NAME' and 'AVG_SALARY'.",
                 expected_query="""
                     SELECT D.NAME AS DEPARTMENT_NAME, AVG(E.SALARY) AS AVG_SALARY
                     FROM EMPLOYEES E
@@ -27,8 +33,10 @@ class EnvironmentTasks:
                 """
             ),
             SqlTask(
+                task_id="senior_engineering_comp_review",
                 difficulty=3,
-                instruction="HARD: Find the name and title of all Engineering (DEPARTMENT_ID = 2) employees who earn a 'Senior' level salary (based on ROLES MIN_SALARY limit).",
+                difficulty_label="hard",
+                instruction="HARD: For compensation review, find Engineering employees whose salary qualifies for the 'Senior' role threshold. Return employee NAME and role TITLE.",
                 expected_query="""
                     SELECT E.NAME, R.TITLE 
                     FROM EMPLOYEES E

inference.py

@@ -1,7 +1,7 @@
 import os
 import textwrap
 from openai import OpenAI
-from typing import List
+from typing import List, Dict, Any
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -15,10 +15,10 @@ ENV_IMAGE_NAME = "sql-agent-env:latest"
 
 # LLM inference config as per instructions
 API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY")
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+API_KEY = os.getenv("OPENAI_API_KEY")
 MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
-MAX_STEPS = 10
+MAX_STEPS = int(os.getenv("MAX_STEPS", "20"))
+OPENAI_SEED = int(os.getenv("OPENAI_SEED", "42"))
 
 SYSTEM_PROMPT = textwrap.dedent(
     """
@@ -64,22 +64,10 @@ def parse_model_action(response_text: str) -> str:
     return query.strip()
 
 def main():
-    if not API_KEY and not GEMINI_API_KEY:
-        print("WARNING: API credentials might be missing. Using dummy environment flow for validation.")
-    
-    # Client setup
-    client_type = "openai"
-    if GEMINI_API_KEY:
-        try:
-            from google import genai
-            client = genai.Client(api_key=GEMINI_API_KEY)
-            client_type = "gemini"
-            print("--- Using Google Gemini Client ---")
-        except ImportError:
-            print("WARNING: google-genai package not found. Falling back to OpenAI client.")
-            client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY or "dummy-key")
-    else:
-        client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY or "dummy-key")
+    if not API_KEY:
+        raise RuntimeError("OPENAI_API_KEY is required for reproducible baseline runs.")
+
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
 
     print("--- SQL Data Analyst OpenEnv Baseline ---")
     
@@ -105,7 +93,9 @@ def main():
                 pass
         env = DirectClient(base_env)
 
-    history = []
+    history: List[str] = []
+    task_scores: Dict[str, float] = {}
+    run_trace: List[Dict[str, Any]] = []
     
     try:
         result = env.reset()
@@ -124,52 +114,59 @@ def main():
                 {"role": "user", "content": user_prompt},
             ]
             
-            try:
-                if client_type == "gemini":
-                    # For Gemini, we collapse System & User prompts into a single payload
-                    gemini_prompt = f"{SYSTEM_PROMPT}\n\n{user_prompt}"
-                    # Use provided model Name, fallback defaults to gemini-2.5-flash
-                    gemini_model = MODEL_NAME if MODEL_NAME != "gpt-4o-mini" else "gemini-2.5-flash"
-                    
-                    response = client.models.generate_content(
-                        model=gemini_model,
-                        contents=gemini_prompt,
-                        config={"temperature": 0.1}
-                    )
-                    response_text = response.text or "SELECT * FROM DEPARTMENTS"
-                elif API_KEY or os.getenv("OPENAI_API_KEY"):
-                    completion = client.chat.completions.create(
-                        model=MODEL_NAME,
-                        messages=messages,
-                        temperature=0.1,
-                        stream=False
-                    )
-                    response_text = completion.choices[0].message.content or "SELECT * FROM DEPARTMENTS"
-                else:
-                    response_text = "SELECT * FROM DEPARTMENTS"  # Dummy mode
-
-            except Exception as e:
-                print(f"Model Request failed: {e}. Using fallback action.")
-                response_text = "SELECT 1"
+            completion = client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=messages,
+                temperature=0.0,
+                seed=OPENAI_SEED,
+                stream=False,
+            )
+            response_text = completion.choices[0].message.content or "SELECT 1"
 
             action_str = parse_model_action(response_text)
             print(f"\n[Step {step}] Model suggested: {action_str}")
-            
+            current_task_id = observation.task_id
             result = env.step(SqlAction(query=action_str))
             observation = result.observation
             reward = result.reward
             done = result.done
 
             print(f"  Reward: {reward:+.2f} | Done: {done}")
+            print(f"  Task Score: {observation.task_score:.3f} | Difficulty: {observation.difficulty}")
+            if observation.grader_feedback:
+                print(f"  Grader: {observation.grader_feedback}")
             if observation.execution_error:
                 print(f"  Error: {observation.execution_error[:200]}")
+
+            if observation.task_score > 0:
+                task_scores[current_task_id] = max(task_scores.get(current_task_id, 0.0), observation.task_score)
                 
             history.append(f"Q: {action_str} -> R: {reward:+.2f}")
+            run_trace.append(
+                {
+                    "step": step,
+                    "task_id": current_task_id,
+                    "action": action_str,
+                    "reward": reward,
+                    "task_score": observation.task_score,
+                    "done": done,
+                }
+            )
             
             if done:
                 print("\nAll tasks complete!")
                 break
-                
+
+        print("\n=== Baseline Summary ===")
+        for task_id in sorted(task_scores):
+            print(f"{task_id}: {task_scores[task_id]:.3f}")
+        task_avg = (sum(task_scores.values()) / len(task_scores)) if task_scores else 0.0
+        total_return = sum(item["reward"] for item in run_trace)
+        print(f"average_task_score: {task_avg:.3f}")
+        print(f"episode_return: {total_return:.3f}")
+        print(f"model: {MODEL_NAME} | seed: {OPENAI_SEED}")
+        print("reproducibility_note: deterministic prompt, temperature=0.0, fixed seed")
+
     finally:
         env.close()
 

models.py

@@ -1,5 +1,5 @@
-from pydantic import Field
-from typing import Optional, List
+from pydantic import BaseModel, Field
+from typing import Optional, List, Dict, Any
 from openenv.core.env_server import Action, Observation, State
 
 class SqlAction(Action):
@@ -10,8 +10,24 @@ class SqlObservation(Observation):
     """Observation returned after executing a SQL query (or at reset)."""
     current_task_instruction: str
     schema_info: str
+    task_id: str
+    difficulty: str
     execution_result: Optional[str] = None
     execution_error: Optional[str] = None
+    task_score: float = 0.0
+    grader_feedback: Optional[str] = None
+
+
+class SqlReward(BaseModel):
+    """Structured reward payload for deterministic scoring details."""
+    total: float
+    task_score: float
+    valid_sql_bonus: float = 0.0
+    progress_bonus: float = 0.0
+    step_penalty: float = 0.0
+    safety_penalty: float = 0.0
+    error_penalty: float = 0.0
+    details: Dict[str, Any] = Field(default_factory=dict)
 
 class SqlState(State):
     """State tracking for the SQL agent environment."""
@@ -19,3 +35,7 @@ class SqlState(State):
     total_tasks: int = 3
     accumulated_reward: float = 0.0
     task_scores: List[float] = Field(default_factory=list)
+    attempts_per_task: List[int] = Field(default_factory=list)
+    max_attempts_per_task: int = 6
+    last_reward: float = 0.0
+    last_info: Dict[str, Any] = Field(default_factory=dict)

pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sqlagent-openenv"
+version = "0.1.0"
+description = "OpenEnv SQL data analyst environment for real-world agent evaluation"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+  "fastapi>=0.111.0",
+  "pydantic>=2.7.1",
+  "uvicorn>=0.30.1",
+  "openenv-core",
+  "pandas",
+  "openai",
+  "python-dotenv",
+]
+
+[project.scripts]
+server = "server.app:main"

requirements.txt

@@ -3,5 +3,5 @@ pydantic>=2.7.1
 uvicorn>=0.30.1
 openenv-core
 pandas
-google-genai
+openai
 python-dotenv

server/app.py

@@ -1,4 +1,7 @@
 from openenv.core.env_server import create_fastapi_app
+from fastapi.responses import HTMLResponse
+import os
+import uvicorn
 
 from models import SqlAction, SqlObservation
 from server.environment import SqlEnvironment
@@ -8,9 +11,33 @@ app = create_fastapi_app(SqlEnvironment, SqlAction, SqlObservation)
 
 @app.get("/")
 def root():
-    return {"status": "ok"}
+    return HTMLResponse(
+        """
+        <html>
+          <head><title>SQL OpenEnv Space</title></head>
+          <body style="font-family: sans-serif; margin: 2rem;">
+            <h2>SQL Data Analyst OpenEnv Environment</h2>
+            <p>Environment is running.</p>
+            <ul>
+              <li>Health check: <a href="/health">/health</a></li>
+              <li>OpenAPI docs: <a href="/docs">/docs</a></li>
+            </ul>
+            <p>Use the OpenEnv client in <code>inference.py</code> to run baseline episodes.</p>
+          </body>
+        </html>
+        """
+    )
 
 
 @app.get("/health")
 def health():
     return {"status": "healthy"}
+
+
+def main():
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run("server.app:app", host="0.0.0.0", port=port)
+
+
+if __name__ == "__main__":
+    main()

server/environment.py

@@ -1,9 +1,8 @@
 import uuid
-import pandas as pd
 from typing import Tuple
 
 from openenv.core.env_server import Environment
-from models import SqlAction, SqlObservation, SqlState
+from models import SqlAction, SqlObservation, SqlState, SqlReward
 from database import SqliteManager
 from core import EnvironmentTasks, DataFrameGrader
 
@@ -17,6 +16,7 @@ class SqlEnvironment(Environment):
         # Instantiate db connection immediately for the server loop
         self.db = SqliteManager()
         self.db.connect()
+        self._dangerous_keywords = ("drop ", "delete ", "truncate ", "alter ", "update ", "insert ")
 
     def reset(self) -> SqlObservation:
         self._state = SqlState(
@@ -24,12 +24,15 @@ class SqlEnvironment(Environment):
             current_task_index=0,
             total_tasks=self.task_registry.get_total_tasks(),
             accumulated_reward=0.0,
-            task_scores=[]
+            task_scores=[],
+            attempts_per_task=[0 for _ in range(self.task_registry.get_total_tasks())],
         )
         first_task = self.task_registry.get_task(0)
         return SqlObservation(
             current_task_instruction=first_task.instruction,
-            schema_info=self.db.get_schema_summary()
+            schema_info=self.db.get_schema_summary(),
+            task_id=first_task.task_id,
+            difficulty=first_task.difficulty_label,
         )
 
     def state(self) -> SqlState:
@@ -37,16 +40,31 @@ class SqlEnvironment(Environment):
 
     def step(self, action: SqlAction) -> Tuple[SqlObservation, float, bool]:
         self._state.step_count += 1
-        query = action.query
+        query = action.query.strip()
         current_idx = self._state.current_task_index
         task = self.task_registry.get_task(current_idx)
+        self._state.attempts_per_task[current_idx] += 1
         
         execution_result_str = None
         error_str = None
-        reward = 0.0
+        score = 0.0
+        grader_feedback = ""
+        reward = SqlReward(total=0.0, task_score=0.0)
         done = False
+
+        query_lower = query.lower()
+        if any(keyword in query_lower for keyword in self._dangerous_keywords):
+            reward.safety_penalty = -1.0
+            reward.total += reward.safety_penalty
+            error_str = "Destructive SQL command detected. Only read-only SELECT queries are allowed."
+        else:
+            reward.step_penalty = -0.02
+            reward.total += reward.step_penalty
         
         try:
+            if error_str:
+                raise ValueError(error_str)
+
             # 1. Gather expected vs agent dataframes
             expected_df = self.db.execute_dataframe(task.expected_query)
             agent_df = self.db.execute_dataframe(query)
@@ -55,41 +73,68 @@ class SqlEnvironment(Environment):
             execution_result_str = agent_df.to_json(orient="records")
             
             # 2. Grade execution
-            score = DataFrameGrader.grade(agent_df, expected_df)
-            self._state.task_scores.append(score)
+            grade_result = DataFrameGrader.grade_with_details(agent_df, expected_df)
+            score = grade_result["score"]
+            grader_feedback = grade_result["feedback"]
+            reward.details = grade_result["details"]
             
-            # 3. Reward scaling
-            reward += 0.1 # Valid execution baseline
-            reward += score
+            # 3. Reward shaping
+            reward.task_score = score
+            reward.valid_sql_bonus = 0.1
+            reward.progress_bonus = 0.9 * score
+            reward.total += reward.valid_sql_bonus + reward.progress_bonus
             
             # 4. Check level progression
-            if score == 1.0:
-                 self._state.current_task_index += 1
-                 if self._state.current_task_index >= self.task_registry.get_total_tasks():
-                     done = True
+            if score >= 0.95:
+                self._state.task_scores.append(score)
+                self._state.current_task_index += 1
+                if self._state.current_task_index >= self.task_registry.get_total_tasks():
+                    done = True
+            elif self._state.attempts_per_task[current_idx] >= self._state.max_attempts_per_task:
+                # Episode ends if agent gets stuck on same task
+                self._state.task_scores.append(score)
+                done = True
                      
         except Exception as e:
             error_str = str(e)
-            reward -= 0.5  # Heavy penalty for destructive/invalid SQL
+            reward.error_penalty = -0.5
+            reward.total += reward.error_penalty
         
         # Global accumulation
-        self._state.accumulated_reward += reward
+        self._state.last_reward = reward.total
+        reward_breakdown = reward.model_dump() if hasattr(reward, "model_dump") else reward.dict()
+        self._state.last_info = {
+            "task_id": task.task_id,
+            "difficulty": task.difficulty_label,
+            "task_score": score,
+            "attempts_for_task": self._state.attempts_per_task[current_idx],
+            "reward_breakdown": reward_breakdown,
+        }
+        self._state.accumulated_reward += reward.total
         
         # Prepare observation response
         if done:
             total_score = sum(self._state.task_scores)
             next_instruction = f"All tasks completed! Final SQL Capability Score: {total_score}"
             schema = ""
+            task_id = task.task_id
+            difficulty = task.difficulty_label
         else:
             next_task = self.task_registry.get_task(self._state.current_task_index)
             next_instruction = next_task.instruction
             schema = self.db.get_schema_summary()
+            task_id = next_task.task_id
+            difficulty = next_task.difficulty_label
             
         obs = SqlObservation(
             current_task_instruction=next_instruction,
             schema_info=schema,
+            task_id=task_id,
+            difficulty=difficulty,
             execution_result=execution_result_str,
-            execution_error=error_str
+            execution_error=error_str,
+            task_score=score,
+            grader_feedback=grader_feedback,
         )
         
-        return obs, reward, done
+        return obs, reward.total, done

uv.lock

@@ -0,0 +1,2 @@
+# Placeholder lockfile for OpenEnv validator compatibility.
+# Generate with `uv lock` in environments where uv is installed.