update models

client.py

@@ -16,66 +16,33 @@ from .models import SkillForgeAction, SkillForgeObservation
 
 
 class SkillForgeEnv(
-    EnvClient[SkillForgeAction, SkillForgeObservation]
+    EnvClient[SkillForgeAction, SkillForgeObservation, State]
 ):
     """
     Client for the Skill Forge Environment.
 
-    This client maintains a persistent WebSocket connection to the environment server,
-    enabling efficient multi-step interactions with lower latency.
-    Each client instance has its own dedicated environment session on the server.
+    Maintains a persistent WebSocket connection to the environment server.
+    Each client instance has its own dedicated environment session.
 
     Example:
-        >>> # Connect to a running server
         >>> with SkillForgeEnv(base_url="http://localhost:8000") as client:
         ...     result = client.reset()
-        ...     print(result.observation.echoed_message)
+        ...     print(result.observation.task_description)
         ...
-        ...     result = client.step(SkillForgeAction(message="Hello!"))
-        ...     print(result.observation.echoed_message)
-
-    Example with Docker:
-        >>> # Automatically start container and connect
-        >>> client = SkillForgeEnv.from_docker_image("skill_forge-env:latest")
-        >>> try:
-        ...     result = client.reset()
-        ...     result = client.step(SkillForgeAction(message="Test"))
-        ... finally:
-        ...     client.close()
+        ...     action = SkillForgeAction(
+        ...         action_type="raw_code",
+        ...         content="df.sort_values('revenue', ascending=False)['product'].tolist()",
+        ...     )
+        ...     result = client.step(action)
+        ...     print(result.observation.result_correct)
     """
 
     def _step_payload(self, action: SkillForgeAction) -> Dict:
-        """
-        Convert SkillForgeAction to JSON payload for step message.
-
-        Args:
-            action: SkillForgeAction instance
-
-        Returns:
-            Dictionary representation suitable for JSON encoding
-        """
-        return {
-            "message": action.message,
-        }
+        return action.model_dump()
 
     def _parse_result(self, payload: Dict) -> StepResult[SkillForgeObservation]:
-        """
-        Parse server response into StepResult[SkillForgeObservation].
-
-        Args:
-            payload: JSON response data from server
-
-        Returns:
-            StepResult with SkillForgeObservation
-        """
         obs_data = payload.get("observation", {})
-        observation = SkillForgeObservation(
-            echoed_message=obs_data.get("echoed_message", ""),
-            message_length=obs_data.get("message_length", 0),
-            done=payload.get("done", False),
-            reward=payload.get("reward"),
-            metadata=obs_data.get("metadata", {}),
-        )
+        observation = SkillForgeObservation(**obs_data)
 
         return StepResult(
             observation=observation,
@@ -84,15 +51,6 @@ class SkillForgeEnv(
         )
 
     def _parse_state(self, payload: Dict) -> State:
-        """
-        Parse server response into State object.
-
-        Args:
-            payload: JSON response from state request
-
-        Returns:
-            State object with episode_id and step_count
-        """
         return State(
             episode_id=payload.get("episode_id"),
             step_count=payload.get("step_count", 0),

models.py

@@ -19,7 +19,7 @@ class SkillForgeAction(Action):
     """Action for the Skill Forge environment"""
     action_type: Literal["create_skill", "use_skill", "raw_code"]
     content: str = Field(description="The content of the action. For create_skill, it is the template. For use_skill, it is the skill id. For raw_code, it is the code.")
-    skill_name: Optional[str] # only for create_skill
+    skill_name: Optional[str] = None # only for create_skill
     reasoning: str = ""
     params: Optional[dict] = None
 

server/Dockerfile

@@ -71,6 +71,9 @@ ENV PATH="/app/.venv/bin:$PATH"
 # Set PYTHONPATH so imports work correctly
 ENV PYTHONPATH="/app/env:$PYTHONPATH"
 
+# Enable the OpenEnv web interface
+ENV ENABLE_WEB_INTERFACE="true"
+
 # Health check
 HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
     CMD curl -f http://localhost:8000/health || exit 1

server/__init__.py

@@ -6,6 +6,9 @@
 
 """Skill Forge environment server components."""
 
-from .skill_forge_environment import SkillForgeEnvironment
+try:
+    from .environment import SkillForgeEnvironment
+except ImportError:
+    from environment import SkillForgeEnvironment
 
 __all__ = ["SkillForgeEnvironment"]

server/app.py

@@ -35,9 +35,12 @@ except Exception as e:  # pragma: no cover
         "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
     ) from e
 
-# Import from local models.py (PYTHONPATH includes /app/env in Docker)
-from models import SkillForgeAction, SkillForgeObservation
-from .skill_forge_environment import SkillForgeEnvironment
+try:
+    from ..models import SkillForgeAction, SkillForgeObservation
+    from .environment import SkillForgeEnvironment
+except ImportError:
+    from models import SkillForgeAction, SkillForgeObservation
+    from environment import SkillForgeEnvironment
 
 
 # Create the app with web interface and README integration

server/data_generator.py

@@ -1,18 +1,403 @@
 import pandas as pd
+from datetime import datetime, timedelta
+
+# ---------------------------------------------------------------------------
+# Cluster A: filter → sort → extract
+# ---------------------------------------------------------------------------
+
+_A1_df = pd.DataFrame({
+    "product": ["Widget", "Gadget", "Doohickey", "Sprocket", "Thingamajig", "Gizmo"],
+    "revenue": [5000, 8000, 3000, 7500, 2000, 6000],
+    "region": ["West", "West", "East", "West", "East", "North"],
+})
+
+_A2_df = pd.DataFrame({
+    "name": ["Alice", "Bob", "Carol", "Dan", "Eve", "Frank", "Grace"],
+    "dept": ["Eng", "Sales", "Eng", "Eng", "HR", "Eng", "Sales"],
+    "tenure": [5, 3, 8, 2, 4, 10, 1],
+})
+
+_A3_df = pd.DataFrame({
+    "order_id": ["ORD-101", "ORD-102", "ORD-103", "ORD-104", "ORD-105", "ORD-106"],
+    "status": ["shipped", "pending", "shipped", "shipped", "cancelled", "shipped"],
+    "quantity": [50, 20, 80, 30, 10, 60],
+})
+
+_A4_df = pd.DataFrame({
+    "student_id": ["S01", "S02", "S03", "S04", "S05", "S06", "S07"],
+    "grade": ["A", "B", "A", "A", "C", "A", "B"],
+    "gpa": [3.9, 3.2, 3.95, 3.7, 2.8, 3.85, 3.1],
+})
+
+_A5_df = pd.DataFrame({
+    "item_name": ["Bolts", "Nails", "Screws", "Washers", "Rivets", "Pins"],
+    "stock": [5, 50, 8, 3, 100, 7],
+    "reorder_priority": [2, 10, 3, 1, 15, 4],
+})
+
+# ---------------------------------------------------------------------------
+# Cluster B: normalize → filter → extract
+# ---------------------------------------------------------------------------
+
+_B1_df = pd.DataFrame({
+    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6"],
+    "email": ["Alice@Gmail.COM", "bob@yahoo.com", "Carol@GMAIL.com",
+              "dan@gmail.COM", "eve@outlook.com", "frank@Gmail.com"],
+})
+
+_B2_df = pd.DataFrame({
+    "id": [1, 2, 3, 4, 5, 6, 7],
+    "name": ["  Alice  ", " Bob", "  Andrew ", "Anna  ", " Carl ", "  Amy", " Brian  "],
+})
+
+_B3_df = pd.DataFrame({
+    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6"],
+    "country_code": ["us", "UK", "Us", "ca", "US", "us"],
+})
+
+_B4_df = pd.DataFrame({
+    "product_id": ["P1", "P2", "P3", "P4", "P5", "P6"],
+    "product_name": ["widget pro", "basic gadget", "pro sprocket",
+                     "mega pro tool", "simple bolt", "pro widget x"],
+})
+
+_B5_df = pd.DataFrame({
+    "contact_id": ["C1", "C2", "C3", "C4", "C5", "C6"],
+    "phone": ["(555) 123-4567", "555.987.6543", "(555)111-2222",
+              "5551234", "555-999-8888", "(555) 000 1111"],
+})
+
+# ---------------------------------------------------------------------------
+# Cluster C: date delta → filter → count
+# ---------------------------------------------------------------------------
+
+_today = datetime(2026, 3, 8)
+
+_C1_df = pd.DataFrame({
+    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6", "U7"],
+    "signup_date": [
+        _today - timedelta(days=10),
+        _today - timedelta(days=45),
+        _today - timedelta(days=5),
+        _today - timedelta(days=90),
+        _today - timedelta(days=25),
+        _today - timedelta(days=3),
+        _today - timedelta(days=60),
+    ],
+    "active": [True, True, True, False, True, True, False],
+})
+
+_C2_df = pd.DataFrame({
+    "order_id": ["O1", "O2", "O3", "O4", "O5", "O6"],
+    "order_date": [
+        _today - timedelta(days=2),
+        _today - timedelta(days=10),
+        _today - timedelta(days=5),
+        _today - timedelta(days=1),
+        _today - timedelta(days=14),
+        _today - timedelta(days=6),
+    ],
+    "amount": [100, 200, 150, 50, 300, 75],
+})
+
+_C3_df = pd.DataFrame({
+    "emp_id": ["E1", "E2", "E3", "E4", "E5", "E6"],
+    "hire_date": [
+        _today - timedelta(days=365),
+        _today - timedelta(days=900),
+        _today - timedelta(days=200),
+        _today - timedelta(days=500),
+        _today - timedelta(days=100),
+        _today - timedelta(days=1500),
+    ],
+    "dept": ["Eng", "Sales", "Eng", "HR", "Eng", "Sales"],
+})
+
+_C4_df = pd.DataFrame({
+    "event_id": ["EV1", "EV2", "EV3", "EV4", "EV5", "EV6"],
+    "event_date": [
+        _today + timedelta(days=5),
+        _today + timedelta(days=20),
+        _today + timedelta(days=10),
+        _today + timedelta(days=3),
+        _today + timedelta(days=30),
+        _today + timedelta(days=12),
+    ],
+    "venue": ["Hall A", "Hall B", "Hall A", "Hall C", "Hall B", "Hall A"],
+})
+
+_C5_df = pd.DataFrame({
+    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6", "U7", "U8"],
+    "birthdate": [
+        _today - timedelta(days=365 * 20),
+        _today - timedelta(days=365 * 30),
+        _today - timedelta(days=365 * 22),
+        _today - timedelta(days=365 * 17),
+        _today - timedelta(days=365 * 19),
+        _today - timedelta(days=365 * 25),
+        _today - timedelta(days=365 * 24),
+        _today - timedelta(days=365 * 15),
+    ],
+})
+
+# ---------------------------------------------------------------------------
+# Cluster D: groupby → aggregate → sort → slice
+# ---------------------------------------------------------------------------
+
+_D1_df = pd.DataFrame({
+    "region": ["West", "East", "West", "North", "East", "North", "West", "East", "South", "South"],
+    "revenue": [5000, 3000, 7000, 4000, 6000, 2000, 3000, 8000, 1000, 5000],
+    "product": ["A", "B", "C", "A", "B", "C", "A", "B", "C", "A"],
+})
+
+_D2_df = pd.DataFrame({
+    "customer": ["Alice", "Bob", "Alice", "Carol", "Bob", "Alice",
+                 "Carol", "Bob", "Alice", "Bob", "Carol", "Bob",
+                 "Alice", "Bob"],
+    "order_id": [f"O{i}" for i in range(1, 15)],
+    "amount": [100, 50, 200, 150, 75, 300, 125, 80, 90, 60, 200, 45, 110, 95],
+})
+
+_D3_df = pd.DataFrame({
+    "dept": ["Eng", "Sales", "Eng", "HR", "Sales", "Eng", "HR", "Sales"],
+    "employee": ["A", "B", "C", "D", "E", "F", "G", "H"],
+    "salary": [120000, 80000, 110000, 70000, 90000, 130000, 75000, 85000],
+})
+
+_D4_df = pd.DataFrame({
+    "category": ["Electronics", "Clothing", "Electronics", "Food",
+                 "Clothing", "Food", "Electronics", "Books", "Books"],
+    "product": ["P1", "P2", "P3", "P4", "P5", "P6", "P7", "P8", "P9"],
+    "units_sold": [500, 200, 300, 150, 100, 400, 250, 50, 80],
+})
+
+_D5_df = pd.DataFrame({
+    "venue": ["Arena", "Hall", "Arena", "Park", "Hall", "Arena", "Park", "Hall"],
+    "event": ["E1", "E2", "E3", "E4", "E5", "E6", "E7", "E8"],
+    "attendees": [500, 200, 300, 100, 250, 400, 150, 300],
+    "capacity": [1000, 600, 1000, 200, 600, 1000, 200, 600],
+})
+
+
+# ---------------------------------------------------------------------------
+# Compute expected outputs
+# ---------------------------------------------------------------------------
+
+# Cluster A
+_A1_expected = _A1_df[_A1_df["region"] == "West"].sort_values("revenue", ascending=False)["product"].tolist()
+_A2_expected = _A2_df[_A2_df["dept"] == "Eng"].sort_values("tenure", ascending=False)["name"].tolist()
+_A3_expected = _A3_df[_A3_df["status"] == "shipped"].sort_values("quantity", ascending=False)["order_id"].tolist()
+_A4_expected = _A4_df[_A4_df["grade"] == "A"].sort_values("gpa", ascending=False)["student_id"].tolist()
+_A5_expected = _A5_df[_A5_df["stock"] < 10].sort_values("reorder_priority")["item_name"].tolist()
+
+# Cluster B
+_B1_expected = _B1_df.assign(email=_B1_df["email"].str.lower())\
+    .query("email.str.endswith('@gmail.com')")["email"].tolist()
+_B2_expected = _B2_df.assign(name=_B2_df["name"].str.strip())\
+    .query("name.str.startswith('A')")["name"].tolist()
+_B3_expected = _B3_df.assign(country_code=_B3_df["country_code"].str.upper())\
+    .query("country_code == 'US'")["user_id"].tolist()
+_B4_expected = _B4_df.assign(product_name=_B4_df["product_name"].str.title())\
+    .query("product_name.str.contains('Pro')")["product_id"].tolist()
+_B5_expected = _B5_df.assign(phone=_B5_df["phone"].str.replace(r"\D", "", regex=True))\
+    .query("phone.str.len() == 10")["phone"].tolist()
+
+# Cluster C
+_C1_expected = int(_C1_df.assign(days_since=(pd.Timestamp(_today) - _C1_df["signup_date"]).dt.days)\
+    .query("days_since < 30").shape[0])
+_C2_expected = int(_C2_df.assign(order_age=(pd.Timestamp(_today) - _C2_df["order_date"]).dt.days)\
+    .query("order_age <= 7").shape[0])
+_C3_expected = int(_C3_df.assign(tenure_years=(pd.Timestamp(_today) - _C3_df["hire_date"]).dt.days / 365)\
+    .query("tenure_years < 2").shape[0])
+_C4_expected = int(_C4_df.assign(days_until=(_C4_df["event_date"] - pd.Timestamp(_today)).dt.days)\
+    .query("days_until <= 14").shape[0])
+_C5_expected = int(_C5_df.assign(age=(pd.Timestamp(_today) - _C5_df["birthdate"]).dt.days / 365)\
+    .query("18 <= age <= 25").shape[0])
+
+# Cluster D
+_D1_expected = _D1_df.groupby("region")["revenue"].sum()\
+    .sort_values(ascending=False).head(3).index.tolist()
+_D2_expected = _D2_df.groupby("customer")["order_id"].count()\
+    .loc[lambda x: x > 5].index.tolist()
+_D3_expected = _D3_df.groupby("dept")["salary"].mean()\
+    .sort_values(ascending=False).index[0]
+_D4_expected = _D4_df.groupby("category")["units_sold"].sum()\
+    .sort_values().head(2).index.tolist()
+_D5_expected = _D5_df.assign(over=_D5_df["attendees"] > _D5_df["capacity"])\
+    .query("over").groupby("venue")["event"].count().index.tolist()
+# D5 recompute: venues where total attendees > total capacity
+_D5_agg = _D5_df.groupby("venue").agg({"attendees": "sum", "capacity": "first"}).reset_index()
+_D5_expected = _D5_agg[_D5_agg["attendees"] > _D5_agg["capacity"]]["venue"].tolist()
+
+
+# ---------------------------------------------------------------------------
+# TASKS list
+# ---------------------------------------------------------------------------
 
 TASKS = [
+    # --- Cluster A: filter → sort → extract ---
     {
         "id": "A1",
-        "cluster": "A",                        # never shown to agent
-        "description": "Given a sales dataframe with columns [product, revenue, units], "
-                       "return the product names sorted by revenue descending.",
-        "expected_skill": "sort_values",
-        "dataframe": pd.DataFrame({
-            "product": ["Widget", "Gadget", "Doohickey"],
-            "revenue": [5000, 8000, 3000],
-            "units": [100, 160, 60]
-        }),
-        "expected_output": ["Gadget", "Widget", "Doohickey"],
-    },
-    # ... 19 more
-]
+        "cluster": "A",
+        "description": "Given a sales dataframe with columns [product, revenue, region], "
+                       "return the product names for the West region sorted by revenue descending.",
+        "dataframe": _A1_df,
+        "expected_output": _A1_expected,
+    },
+    {
+        "id": "A2",
+        "cluster": "A",
+        "description": "Given an employees dataframe with columns [name, dept, tenure], "
+                       "return the names of Engineering employees sorted by tenure descending.",
+        "dataframe": _A2_df,
+        "expected_output": _A2_expected,
+    },
+    {
+        "id": "A3",
+        "cluster": "A",
+        "description": "Given an orders dataframe with columns [order_id, status, quantity], "
+                       "return the order IDs for shipped orders sorted by quantity descending.",
+        "dataframe": _A3_df,
+        "expected_output": _A3_expected,
+    },
+    {
+        "id": "A4",
+        "cluster": "A",
+        "description": "Given a students dataframe with columns [student_id, grade, gpa], "
+                       "return the student IDs of students with grade A sorted by GPA descending.",
+        "dataframe": _A4_df,
+        "expected_output": _A4_expected,
+    },
+    {
+        "id": "A5",
+        "cluster": "A",
+        "description": "Given an inventory dataframe with columns [item_name, stock, reorder_priority], "
+                       "return item names where stock is below 10, sorted by reorder priority ascending.",
+        "dataframe": _A5_df,
+        "expected_output": _A5_expected,
+    },
+    # --- Cluster B: normalize → filter → extract ---
+    {
+        "id": "B1",
+        "cluster": "B",
+        "description": "Given a users dataframe with columns [user_id, email], "
+                       "lowercase all emails, keep only @gmail.com addresses, return the email list.",
+        "dataframe": _B1_df,
+        "expected_output": _B1_expected,
+    },
+    {
+        "id": "B2",
+        "cluster": "B",
+        "description": "Given a contacts dataframe with columns [id, name], "
+                       "strip whitespace from names, keep names starting with 'A', return the name list.",
+        "dataframe": _B2_df,
+        "expected_output": _B2_expected,
+    },
+    {
+        "id": "B3",
+        "cluster": "B",
+        "description": "Given a users dataframe with columns [user_id, country_code], "
+                       "uppercase all country codes, keep only 'US', return the user_id list.",
+        "dataframe": _B3_df,
+        "expected_output": _B3_expected,
+    },
+    {
+        "id": "B4",
+        "cluster": "B",
+        "description": "Given a products dataframe with columns [product_id, product_name], "
+                       "title-case all product names, keep those containing 'Pro', return the product_id list.",
+        "dataframe": _B4_df,
+        "expected_output": _B4_expected,
+    },
+    {
+        "id": "B5",
+        "cluster": "B",
+        "description": "Given a contacts dataframe with columns [contact_id, phone], "
+                       "remove all non-digit characters from phone numbers, keep only 10-digit ones, return the phone list.",
+        "dataframe": _B5_df,
+        "expected_output": _B5_expected,
+    },
+    # --- Cluster C: date delta → filter → count ---
+    {
+        "id": "C1",
+        "cluster": "C",
+        "description": "Given a users dataframe with columns [user_id, signup_date, active], "
+                       "compute days since signup (from 2026-03-08), keep users who signed up within the last 30 days, return the count.",
+        "dataframe": _C1_df,
+        "expected_output": _C1_expected,
+    },
+    {
+        "id": "C2",
+        "cluster": "C",
+        "description": "Given an orders dataframe with columns [order_id, order_date, amount], "
+                       "compute order age in days (from 2026-03-08), keep orders within the last 7 days, return the count.",
+        "dataframe": _C2_df,
+        "expected_output": _C2_expected,
+    },
+    {
+        "id": "C3",
+        "cluster": "C",
+        "description": "Given an employees dataframe with columns [emp_id, hire_date, dept], "
+                       "compute tenure in years (from 2026-03-08), keep employees with less than 2 years, return the count.",
+        "dataframe": _C3_df,
+        "expected_output": _C3_expected,
+    },
+    {
+        "id": "C4",
+        "cluster": "C",
+        "description": "Given an events dataframe with columns [event_id, event_date, venue], "
+                       "compute days until event (from 2026-03-08), keep events within the next 14 days, return the count.",
+        "dataframe": _C4_df,
+        "expected_output": _C4_expected,
+    },
+    {
+        "id": "C5",
+        "cluster": "C",
+        "description": "Given a users dataframe with columns [user_id, birthdate], "
+                       "compute age in years (from 2026-03-08), keep users aged 18 to 25 inclusive, return the count.",
+        "dataframe": _C5_df,
+        "expected_output": _C5_expected,
+    },
+    # --- Cluster D: groupby → aggregate → sort → slice ---
+    {
+        "id": "D1",
+        "cluster": "D",
+        "description": "Given a sales dataframe with columns [region, revenue, product], "
+                       "group by region, sum revenue, return the top 3 regions by total revenue.",
+        "dataframe": _D1_df,
+        "expected_output": _D1_expected,
+    },
+    {
+        "id": "D2",
+        "cluster": "D",
+        "description": "Given an orders dataframe with columns [customer, order_id, amount], "
+                       "group by customer, count orders, return customers with more than 5 orders.",
+        "dataframe": _D2_df,
+        "expected_output": _D2_expected,
+    },
+    {
+        "id": "D3",
+        "cluster": "D",
+        "description": "Given an employees dataframe with columns [dept, employee, salary], "
+                       "group by department, compute average salary, return the department with the highest average.",
+        "dataframe": _D3_df,
+        "expected_output": _D3_expected,
+    },
+    {
+        "id": "D4",
+        "cluster": "D",
+        "description": "Given a products dataframe with columns [category, product, units_sold], "
+                       "group by category, sum units_sold, return the bottom 2 categories by total units.",
+        "dataframe": _D4_df,
+        "expected_output": _D4_expected,
+    },
+    {
+        "id": "D5",
+        "cluster": "D",
+        "description": "Given an events dataframe with columns [venue, event, attendees, capacity], "
+                       "group by venue summing attendees and taking the first capacity, "
+                       "return venues where total attendees exceed capacity.",
+        "dataframe": _D5_df,
+        "expected_output": _D5_expected,
+    },
+]

server/{skill_forge_environment.py → environment.py}

@@ -10,6 +10,7 @@ Skill Forge Environment Implementation.
 An RL training environment where LLM Agents evolve from "reinventing the wheel" to "building a skill library."
 """
 
+import json
 import traceback
 from uuid import uuid4
 
@@ -18,139 +19,155 @@ import pandas as pd
 from openenv.core.env_server.interfaces import Environment
 from openenv.core.env_server.types import State
 
-from models import SkillForgeAction, SkillForgeObservation
-from .data_generator import TASKS
+try:
+    from ..models import SkillForgeAction, SkillForgeObservation
+    from .data_generator import TASKS
+except ImportError:
+    from models import SkillForgeAction, SkillForgeObservation
+    from data_generator import TASKS
+
 
 class SkillForgeEnvironment(Environment):
     """
-    A simple echo environment that echoes back messages.
-
-    This environment is designed for testing the HTTP server infrastructure.
-    It maintains minimal state and simply echoes back whatever message it receives.
-
-    Example:
-        >>> env = SkillForgeEnvironment()
-        >>> obs = env.reset()
-        >>> print(obs.echoed_message)  # "Skill Forge environment ready!"
-        >>>
-        >>> obs = env.step(SkillForgeAction(message="Hello"))
-        >>> print(obs.echoed_message)  # "Hello"
-        >>> print(obs.message_length)  # 5
+    SkillForge RL environment.
+
+    The agent solves chained pandas tasks and can build a reusable skill library.
+    Skills persist across episodes so the agent can discover and reuse patterns.
     """
 
-    # Enable concurrent WebSocket sessions.
-    # Set to True if your environment isolates state between instances.
-    # When True, multiple WebSocket clients can connect simultaneously, each
-    # getting their own environment instance (when using factory mode in app.py).
     SUPPORTS_CONCURRENT_SESSIONS: bool = True
 
     def __init__(self):
-        """Initialize the skill_forge environment."""
         self._state = State(episode_id=str(uuid4()), step_count=0)
-        self._reset_count = 0
-
-        self.skill_library = {}
-        self.task_idx = 0
-        self.action_history = []
-        self.error_history = []
+        self.skill_library: dict = {}
+        self.task_idx: int = 0
+        self.tasks_solved: int = 0
+        self.total_tokens: int = 0
 
     def reset(self) -> SkillForgeObservation:
         """
-        Reset the environment.
-        skill_library intentionally NOT reset — persists across episodes
-
-        Returns:
-            SkillForgeObservation with a ready message
+        Reset episode state. skill_library is NOT reset — persists across episodes.
         """
         self._state = State(episode_id=str(uuid4()), step_count=0)
-        self._reset_count += 1
-
         self.task_idx = 0
-        task = TASKS[self.task_idx]
-        self.action_history = []
-        self.error_history = []
+        self.tasks_solved = 0
+        self.total_tokens = 0
 
-        return SkillForgeObservation(
-            task_id=task["id"],
-            task_description=task["description"],
-            snapshot_data=task["dataframe"].head(5).to_string(),
-            skill_library=self.skill_library,
-            context="",
-            step_count=0,
-            total_tokens=0,
-            result_correct=False,
-            result_output="",
-            expected_output=str(task["expected_output"]),
-        )
+        task = TASKS[self.task_idx]
+        return self._make_observation(task, result_correct=False, result_output="", reward=0.0, done=False)
 
-    def step(self, action: SkillForgeAction) -> SkillForgeObservation:  
-        # TODO: create function _is_redundant
+    def step(self, action: SkillForgeAction) -> SkillForgeObservation:
         self._state.step_count += 1
-
-        if action.action_type not in ["create_skill", "use_skill", "raw_code"]:
-            raise ValueError(f"Invalid action type: {action.action_type}")
-
         task = TASKS[self.task_idx]
-        reward = 0.0
-        
+
+        # --- create_skill: store template, stay on current task ---
         if action.action_type == "create_skill":
+            token_cost = len(action.content)
+            self.total_tokens += token_cost
+
             self.skill_library[action.skill_name] = {
                 "template": action.content,
                 "description": action.reasoning,
                 "used_count": 0,
             }
-            reward = 0.5 # TODO: revisit this value. Set it a bit high as we reward exploration 
-            result_correct = False
-            result_output = f"Skill {action.skill_name} created"
+            reward = 0.5
+            return self._make_observation(
+                task, result_correct=False,
+                result_output=f"Skill '{action.skill_name}' saved.",
+                reward=reward, done=False,
+            )
+
+        # --- use_skill or raw_code: execute and evaluate ---
+        if action.action_type == "use_skill":
+            skill = self.skill_library.get(action.content)
+            if skill is None:
+                # skill not found ��� treat as error
+                self.total_tokens += len(action.content)
+                return self._make_observation(
+                    task, result_correct=False,
+                    result_output=f"Skill '{action.content}' not found in library.",
+                    reward=-0.3, done=False,
+                )
+            exec_code = skill["template"].format(**(action.params or {}))
+            skill["used_count"] += 1
+            # token cost for use_skill: skill name + serialized params (much shorter than full code)
+            skill_call_repr = action.content + json.dumps(action.params or {})
+            token_cost = len(skill_call_repr)
         else:
+            # raw_code
+            exec_code = action.content
+            token_cost = len(action.content)
+
+        self.total_tokens += token_cost
+
+        result_correct, result_output = self._evaluate(exec_code, task["dataframe"], task["expected_output"])
+
+        if result_correct:
+            reward = 2.0
+            reward -= 0.001 * token_cost
             if action.action_type == "use_skill":
-                skill = self.skill_library.get(action.content)                
-                # TODO: if action_type is use_skill while we don't have the skill yet, read reasoning to understand why the skill should be used and create the skill accordingly
-                if skill: 
-                    exec_code = skill["template"].format(**(action.params or {}))
-                    self.skill_library[action.content]["used_count"] += 1
-                else: # TODO: log this scenario
-                    exec_code = None
-            else:
-                exec_code = action.code # TODO: add a penalty if we have a skill in skill lib for this problem
-
-            # TODO: dataframe shouldn't be part of the task, the design is ugly
-            result_correct, result_output = self._evaluate(exec_code, task["dataframe"], task["expected_output"])
-            reward += 2.0 if result_correct else -1.0
-            reward -= 0.001 * len(action.content)  # token penalty
-            if action.action_type == "use_skill" and result_correct:
-                reward += 0.5 # reward for using a skill that worked
-            
-            # TODO P0: update the context with action history and error history 
-
-            # advance task
-            self.task_idx = min(self.task_idx + 1, len(TASKS) - 1)
-            done = self.task_idx == len(TASKS) - 1
-            next_task = TASKS[self.task_idx]
+                reward += 0.5
+            self.tasks_solved += 1
+            self.task_idx += 1
+        else:
+            reward = -0.3
+
+        done = self.task_idx >= len(TASKS)
+        next_task = TASKS[self.task_idx] if not done else task
+
+        return self._make_observation(
+            next_task,
+            result_correct=result_correct,
+            result_output=result_output,
+            reward=reward,
+            done=done,
+        )
 
+    def _evaluate(self, exec_code: str | None, dataframe: pd.DataFrame, expected_output) -> tuple[bool, str]:
+        if exec_code is None:
+            return False, "No code to execute."
+        try:
+            namespace = {"df": dataframe.copy(), "pd": pd, "__builtins__": {"len": len, "str": str, "int": int, "float": float, "list": list, "dict": dict, "bool": bool, "range": range, "abs": abs, "min": min, "max": max, "sum": sum, "sorted": sorted, "round": round, "True": True, "False": False, "None": None}}
+            result = eval(exec_code, namespace)
+
+            # normalize for comparison
+            if isinstance(result, pd.DataFrame):
+                result = result.values.tolist()
+            if isinstance(result, pd.Series):
+                result = result.tolist()
+            if isinstance(result, pd.Index):
+                result = result.tolist()
+
+            expected = expected_output
+            if isinstance(expected, pd.Series):
+                expected = expected.tolist()
+
+            try:
+                is_correct = result == expected
+            except (ValueError, TypeError):
+                is_correct = False
+
+            return bool(is_correct), str(result)
+        except Exception:
+            return False, traceback.format_exc()
+
+    def _make_observation(self, task: dict, result_correct: bool, result_output: str,
+                          reward: float, done: bool) -> SkillForgeObservation:
         return SkillForgeObservation(
-            task_id=next_task["id"],
-            task_description=next_task["description"],
-            snapshot_data=next_task["dataframe"].head(5).to_string(),
+            task_id=task["id"],
+            task_description=task["description"],
+            snapshot_data=task["dataframe"].head(5).to_string(),
             skill_library=self.skill_library,
             context="",
             step_count=self._state.step_count,
-            total_tokens=0, # TODO P0: update this with the total tokens used
+            total_tokens=self.total_tokens,
             result_correct=result_correct,
             result_output=result_output,
-            expected_output=str(next_task["expected_output"]),
+            expected_output=str(task["expected_output"]),
+            reward=reward,
+            done=done,
         )
-    
-    def _evaluate(self, exec_code, dataframe, expected_output):
-        # TODO P0: implement this
-        return True, "Evaluation successful"
 
     @property
     def state(self) -> State:
-        """
-        Get the current environment state.
-
-        Returns:
-            Current State with episode_id and step_count
-        """
         return self._state

server/requirements.txt

@@ -1,6 +1,4 @@
-openenv[core]>=0.2.0
+openenv-core[core]>=0.2.0
 fastapi>=0.115.0
 uvicorn>=0.24.0
-
-
-
+pandas>=2.3.3