Tsah00/sql-env

SQL Query Learning Environment

An OpenEnv-compatible reinforcement learning environment that simulates the day-to-day work of a data analyst fulfilling ad-hoc SQL requests from business stakeholders (marketing, finance, CRM, merchandising). Agents learn to translate natural-language business requirements into correct, efficient SQL queries — a high-value real-world task.

Why This Environment Matters

Text-to-SQL is one of the most commercially valuable analyst skills — data teams spend significant time writing bespoke queries for stakeholder requests, and current LLMs still fail on complex queries involving CTEs, window functions, and multi-table joins. Existing benchmarks (Spider, BIRD) are static evaluation sets. There is no standard RL environment for training agents to improve iteratively on SQL tasks through trial-and-error feedback.

This environment fills that gap:

• —Real-world task: agents play the role of a data analyst, fulfilling concrete business requests — not solving academic puzzles
• —Iterative learning: reward signal at every step, not just at episode end
• —Partial credit with exploit protection: Jaccard-based row matching gives meaningful gradient signal while penalising both missing rows and extra rows — returning the entire table does not score 1.0
• —Difficulty curriculum: 3 tiers (easy → medium → hard) enable curriculum learning strategies
• —Business-realistic framing: all 9 tasks are modelled on real analyst/engineer requests with named stakeholder teams
• —Zero external dependencies: pure SQLite, runs on any 2 vCPU / 8 GB machine

Environment Overview

The environment hosts an in-memory SQLite e-commerce database with 4 tables and deterministic seed data:

At each step the agent submits a SQL query. The grader executes it, compares results to the reference solution, and returns a reward in [0.0, 1.0] with detailed feedback.

Action Space

Observation Space

Tasks (9 total across 3 tiers)

All tasks simulate real stakeholder requests submitted to a data analyst:

Easy — Basic filtering and aggregation

Medium — Joins, grouping, and date logic

Hard — CTEs, window functions, correlated subqueries

Reward Function

reward = correctness × 0.7 + keyword_bonus × 0.1 + efficiency_bonus × 0.2

Why this reward design is good for RL:

• —Non-binary: agents receive gradient signal even on partially correct queries
• —Exploit-resistant: Jaccard scoring penalises both over-fetching and under-fetching
• —Column aliases normalised: total instead of total_spent is not penalised if values match
• —Efficiency signal discourages degenerate solutions
• —Per-step rewards enable policy gradient methods without sparse returns

Grader variety (not always the same score):

Baseline Scores

Deterministic rule-based baseline (no LLM) — achieved in 1 step per task:

LLM agents (e.g. Qwen2.5-72B, Nemotron) are expected to score 0.7–1.0 on easy, 0.5–0.9 on medium, and 0.2–0.7 on hard — leaving meaningful room for RL improvement.

Setup

Local Development

pip install -r requirements.txt

# Start the server
uvicorn server.app:app --host 0.0.0.0 --port 7860

# Run inference (separate terminal)
python inference.py

Docker

docker build -t sql_env .
docker run -p 7860:7860 \
  -e API_BASE_URL="https://router.huggingface.co/v1" \
  -e MODEL_NAME="Qwen/Qwen2.5-72B-Instruct" \
  -e HF_TOKEN="your-hf-token" \
  sql_env

Run Tests

python -m pytest test_env.py -v
# 72 tests: database, graders, partial credit, lifecycle, all 9 tasks, FastAPI endpoints

API Endpoints

Environment Variables

The inference script accepts HF_TOKEN, OPENAI_API_KEY, or API_KEY (checked in that order).

Project Structure

.
├── inference.py           # Baseline inference script (mandatory, root)
├── models.py              # Pydantic Action / Observation / State models
├── client.py              # HTTP EnvClient
├── openenv.yaml           # OpenEnv manifest
├── pyproject.toml         # Package metadata and dependencies
├── uv.lock                # uv lockfile for multi-mode deployment
├── requirements.txt       # pip dependencies
├── Dockerfile             # HF Spaces-compatible container
├── test_env.py            # 72-test suite
└── server/
    ├── app.py             # FastAPI server (reset / step / state / ws)
    ├── sql_environment.py # Core environment logic
    └── tasks.py           # 9 tasks, graders, SQLite seed data

Judging Criteria Compliance