Spaces:
Sleeping
Sleeping
| title: Data Analysis Agent Environment | |
| emoji: 📊 | |
| colorFrom: blue | |
| colorTo: indigo | |
| sdk: docker | |
| pinned: false | |
| # Data Analysis Agent Environment | |
| An OpenEnv-compliant RL environment for training and evaluating data analysis agents. Agents execute pandas code against a business dataset to answer analytical questions, graded by deterministic programmatic graders. | |
| ## Motivation | |
| Data analysis is a universal real-world task. Every business needs analysts who can query datasets, compute metrics, and extract insights. This environment lets RL agents practice that exact workflow — explore a dataset with code, then submit a precise answer — with automatic scoring. | |
| ## Action & Observation Spaces | |
| ### Action (`DataAction`) | |
| | Field | Type | Description | | |
| |---|---|---| | |
| | `action_type` | `"execute_code"` or `"submit_answer"` | What the agent wants to do | | |
| | `code` | `str` (optional) | Python/pandas code to execute | | |
| | `answer` | `str` (optional) | Final answer to submit for grading | | |
| ### Observation (`DataObservation`) | |
| | Field | Type | Description | | |
| |---|---|---| | |
| | `output` | `str` | Stdout from code execution or environment messages | | |
| | `success` | `bool` | Whether the action succeeded | | |
| | `error` | `str` (optional) | Error message if action failed | | |
| | `task_description` | `str` | The question to answer (set on reset) | | |
| | `dataset_info` | `str` | Dataset schema summary (set on reset) | | |
| | `done` | `bool` | Whether the episode is over | | |
| | `reward` | `float` | Step reward | | |
| ### State (`DataState`) | |
| | Field | Type | Description | | |
| |---|---|---| | |
| | `episode_id` | `str` | Unique episode identifier | | |
| | `step_count` | `int` | Current step number | | |
| | `task_id` | `int` | Active task (1–6) | | |
| | `answer_submitted` | `bool` | Whether final answer was submitted | | |
| | `final_score` | `float` | Graded score after submission | | |
| ## Tasks | |
| Tasks use two data sources: | |
| - **`df`** — synthetic e-commerce sales CSV (~2000 orders): `order_id`, `customer_id`, `product_name`, `category`, `quantity`, `unit_price`, `total_price`, `order_date`, `city`, `country` | |
| - **SQLite DB** (`store_data.db`) — additional tables for cross-source tasks: `customer_profiles` (300 rows), `product_catalog` (25 rows) | |
| ### Task 1 — Easy: Top Revenue Category | |
| - **Question**: What is the top-selling product category by total revenue? | |
| - **Grading**: Containment match (case-insensitive) → 1.0 or 0.0 | |
| - **Expected difficulty**: Single groupby + sum + argmax | |
| ### Task 2 — Medium: City Revenue Share | |
| - **Question**: Which city generates the most revenue? What percentage of total revenue does it represent? | |
| - **Grading**: 0.5 for correct city + 0.5 for percentage within ±0.1% | |
| - **Expected difficulty**: Groupby + percentage calculation + formatting | |
| ### Task 3 — Medium: Repeat Customer Cohort Analysis | |
| - **Question**: How many unique customers ordered in both January and December? Compare their average order value to all other customers. | |
| - **Grading**: 0.33 per correct field (count, cohort AOV, other AOV) | |
| - **Expected difficulty**: Temporal filtering, set intersection, conditional aggregation | |
| ### Task 4 — Hard: Monthly Revenue Ratio | |
| - **Question**: Which month had the highest vs. lowest total revenue? What is the ratio between them? | |
| - **Grading**: 0.33 for best month + 0.33 for worst month + 0.34 for ratio within ±0.01 | |
| - **Expected difficulty**: Monthly resample/groupby, min/max comparison, ratio formatting | |
| ### Task 5 — Hard: Customer Loyalty Tier Revenue (cross-source) | |
| - **Question**: Which customer loyalty tier generates the highest total revenue and what percentage does it represent? | |
| - **Data**: Requires joining `df` with `customer_profiles` table from SQLite on `customer_id` | |
| - **Grading**: 0.33 for tier name + 0.33 for revenue within ±0.5% + 0.34 for percentage within ±0.1 | |
| - **Expected difficulty**: SQLite query → pandas merge → groupby aggregation | |
| ### Task 6 — Hard: Supplier Profitability (cross-source) | |
| - **Question**: Which supplier has the highest total profit? What is their average profit margin? | |
| - **Data**: Requires joining `df` with `product_catalog` table from SQLite on `product_name` | |
| - **Grading**: 0.33 for supplier name + 0.34 for total profit within ±0.5% + 0.33 for avg margin within ±0.1 | |
| - **Expected difficulty**: SQLite query → pandas merge → per-order profit/margin calculation → group aggregation | |
| ## Reward Function | |
| | Event | Reward | | |
| |---|---| | |
| | Successful code execution | +0.05 | | |
| | Code execution error | -0.05 | | |
| | Final answer (graded) | 0.0 — 1.0 based on task grader | | |
| | Max steps (20) exceeded | 0.0 | | |
| ## Setup & Usage | |
| ### Prerequisites | |
| - Python 3.13+ | |
| - [uv](https://docs.astral.sh/uv/) package manager | |
| ### Install | |
| ```bash | |
| uv sync | |
| ``` | |
| ### Run the server | |
| ```bash | |
| uv run uvicorn server.app:app --host 0.0.0.0 --port 8000 | |
| ``` | |
| ### Run the inference | |
| - First export all the required env variables mentioned in the .env.example. Then run below command | |
| ```bash | |
| uv run python inference.py | |
| ``` | |
| ### Run the baseline | |
| ```bash | |
| OPENAI_API_KEY=sk-... uv run python baseline.py | |
| # Against a deployed HF Space: | |
| OPENAI_API_KEY=sk-... uv run python baseline.py --base-url https://<your-username>-<space-name>.hf.space | |
| ``` | |
| ### Docker (local) | |
| ```bash | |
| docker build -t data-analysis-env . | |
| docker run -p 7860:7860 data-analysis-env | |
| ``` | |
| ### Client usage (Python) | |
| ```python | |
| from client import DataAnalysisClient | |
| from models import DataAction | |
| # Async | |
| async with DataAnalysisClient(base_url="http://localhost:8000") as client: | |
| result = await client.reset(task_id=1) | |
| result = await client.step(DataAction(action_type="execute_code", code="print(df.head())")) | |
| result = await client.step(DataAction(action_type="submit_answer", answer="Electronics")) | |
| # Sync | |
| with DataAnalysisClient(base_url="http://localhost:8000").sync() as client: | |
| result = client.reset(task_id=2) | |
| result = client.step(DataAction(action_type="execute_code", code="print(df.groupby('city')['total_price'].sum())")) | |
| ``` | |
| ## Project Structure | |
| ``` | |
| ├── models.py # DataAction, DataObservation, DataState | |
| ├── client.py # DataAnalysisClient (EnvClient subclass) | |
| ├── inference.py # HF inference script (uses HF Inference API) | |
| ├── baseline.py # OpenAI baseline inference script | |
| ├── helpers/ | |
| │ └── response_parser.py # Robust LLM JSON response parser | |
| ├── tasks/ | |
| │ ├── base_task.py # Task ABC with grade() interface | |
| │ ├── task_easy.py # Task 1 (Easy): Top revenue category | |
| │ ├── task_medium.py # Task 2 (Medium): City revenue share | |
| │ ├── task_medium_2.py # Task 4 (Hard): Monthly revenue ratio | |
| │ ├── task_hard.py # Task 3 (Medium): Repeat customer cohort | |
| │ ├── task_hard_2.py # Task 5 (Hard): Customer loyalty tier revenue | |
| │ └── task_hard_3.py # Task 6 (Hard): Supplier profitability | |
| ├── datasets/ | |
| │ ├── sales.csv # Synthetic e-commerce sales dataset | |
| │ └── store_data.db # SQLite DB: customer_profiles, product_catalog | |
| ├── server/ | |
| │ ├── app.py # FastAPI app entry point | |
| │ └── data_analysis_env.py # Environment implementation | |
| ├── Dockerfile # HF Spaces Docker build (port 7860) | |
| ├── openenv.yaml # OpenEnv spec metadata | |
| └── pyproject.toml # Dependencies and project config | |
| ``` | |