refactor: Improved grading logic for task 2

agents/task2.py

@@ -30,11 +30,8 @@ def oracle_t2(env: Task2Environment, seed: int, verbose: bool = False) -> Dict[s
     env.step(Action(action_type=ActionType.GET_FUNCTION_CODE))
     result = env.step(Action(action_type=ActionType.SUBMIT_PROPERTY,
                               params={"property": gt_text}))
-    r_val  = result.reward.value
-    score  = round(r_val / 5.0, 4) if r_val > 0 else 0.0
     return {"seed": seed, "contract": contract, "function": fn_name,
-            "grader_score": score,
-            "cumulative_reward": result.observation.cumulative_reward}
+            "grader_score": result.reward.value }
 
 
 def partial_t2(env: Task2Environment, seed: int) -> Dict[str, Any]:
@@ -51,9 +48,7 @@ def partial_t2(env: Task2Environment, seed: int) -> Dict[str, Any]:
             break
     result = env.step(Action(action_type=ActionType.SUBMIT_PROPERTY,
                               params={"property": comment}))
-    r_val = result.reward.value
-    return {"seed": seed, "grader_score": round(r_val / 5.0, 4) if r_val > 0 else 0.0,
-            "cumulative_reward": result.observation.cumulative_reward}
+    return {"seed": seed, "grader_score": result.reward.value}
 
 
 def random_t2(env: Task2Environment, seed: int) -> Dict[str, Any]:
@@ -90,10 +85,8 @@ def random_t2(env: Task2Environment, seed: int) -> Dict[str, Any]:
     prop = rng.choice(templates)
     result = env.step(Action(action_type=ActionType.SUBMIT_PROPERTY,
                               params={"property": prop}))
-    r_val = result.reward.value
-    return {"seed": seed, "grader_score": round(r_val / 5.0, 4) if r_val > 0 else 0.0,
-            "submitted": prop[:60],
-            "cumulative_reward": result.observation.cumulative_reward}
+    return {"seed": seed, "grader_score": result.reward.value,
+            "submitted": prop[:60]}
 
 
 def floor_t2(env: Task2Environment, seed: int) -> Dict[str, Any]:
@@ -101,5 +94,4 @@ def floor_t2(env: Task2Environment, seed: int) -> Dict[str, Any]:
     env.reset(seed=seed)
     result = env.step(Action(action_type=ActionType.SUBMIT_PROPERTY,
                               params={"property": ""}))
-    return {"seed": seed, "grader_score": 0.0,
-            "cumulative_reward": result.observation.cumulative_reward}
+    return {"seed": seed, "grader_score": 0.001}

data/contracts.json

@@ -227,7 +227,7 @@
                 "property_specification": {
                     "precondition": "User has AToken balance B",
                     "operation": "burn(user, receiver, amount, index)",
-                    "expected_postcondition": "User's AToken balance = B - amount (within rounding tolerance ε)",
+                    "postcondition": "User's AToken balance = B - amount (within rounding tolerance ε)",
                     "actual": "When amount.rayDiv(index) rounds down to 0, the burn operation transfers amount underlying tokens but burns 0 ATokens, resulting in user AToken balance unchanged = B, violating the postcondition where the balance should be B - amount."
                 }
             },
@@ -934,7 +934,7 @@
                 "property_specification": {
                     "precondition": "User has debt balance B",
                     "operation": "mint(user, onBehalfOf, amount, rate)",
-                    "expected_postcondition": "User's debt balance = B + amount (within rounding tolerance ε)",
+                    "postcondition": "User's debt balance = B + amount (within rounding tolerance ε)",
                     "actual": "When amount conversion rounds down to 0 in intermediate calculations, the mint operation may mint zero debt tokens while still transferring underlying tokens (or vice versa), resulting in user debt balance unchanged = B, violating the postcondition where the balance should be B + amount."
                 }
             },

env/schemas.py

@@ -39,17 +39,17 @@ class ActionType(str, Enum):
     SUBMIT               = ("submit", 0.0)
 
     # ── Task 2 – Property Discovery ─────────────────────────────────────────
-    GET_SIMILAR_RULE      = ("get_similar_rule", 0.0)
-    GET_FILE_NATSPEC      = ("get_file_natspec", 0.0)
-    GET_FUNCTION_NATSPEC  = ("get_function_natspec", 0.0)
-    GET_RELATED_FUNCTIONS = ("get_related_functions", 0.0)
-    GET_SIGNATURE         = ("get_signature", 0.0)
+    GET_SIMILAR_RULE      = ("get_similar_rule", 0.15)
+    GET_FILE_NATSPEC      = ("get_file_natspec", 0.05)
+    GET_FUNCTION_NATSPEC  = ("get_function_natspec", -0.08)
+    GET_RELATED_FUNCTIONS = ("get_related_functions", 0.07)
+    GET_SIGNATURE         = ("get_signature", 0.04)
     SUBMIT_PROPERTY       = ("submit_property", 0.0)
 
     # ── Task 3 – Rule Checker ────────────────────────────────────────────────
-    GET_PROPERTY_SPECIFICATION = ("get_property_specification", 0.0)
-    GET_FUNCTION_METADATA   = ("get_function_metadata", 0.0)
-    SUBMIT_FUNCTION         = ("submit_function", 0.0)
+    GET_PROPERTY_SPECIFICATION  = ("get_property_specification", 0.0)
+    GET_FUNCTION_METADATA       = ("get_function_metadata", 0.0)
+    SUBMIT_FUNCTION             = ("submit_function", 0.0)
     
     # ─────── General Actions ─────────────────────────────────────────────────
     UNKNOWN                = ("unknown", 0.0)

eval.py

@@ -121,7 +121,7 @@ def run_task2_eval(n: int, seed_offset: int, verbose: bool) -> Dict[str, Any]:
         oracle_eps.append(ep)
         icon = "✅" if ep["grader_score"] >= 0.65 else "⚠️ "
         print(f"  {icon} seed={ep['seed']:3d}  {ep['contract']:12s}.{ep['function']:18s}"
-              f"  score={ep['grader_score']:.3f}  reward={ep['cumulative_reward']:+.2f}")
+              f"  score={ep['grader_score']:.3f}")
     oracle_avg = _avg(oracle_eps)
     print(f"\n  Oracle   avg: {oracle_avg:.3f}")
 
@@ -143,10 +143,10 @@ def run_task2_eval(n: int, seed_offset: int, verbose: bool) -> Dict[str, Any]:
     floor_avg = _avg(floor_eps)
     print(f"  Floor    avg: {floor_avg:.3f}")
 
-   # assert oracle_avg > 0.60,  f"Oracle avg {oracle_avg:.3f} should be > 0.60"
-   # assert oracle_avg > partial_avg >= floor_avg, \
-        # "Score ordering violated: oracle > partial >= floor"
-   # assert floor_avg == 0.0, f"Floor avg {floor_avg:.3f} should be 0.0"
+    assert oracle_avg > 0.60,  f"Oracle avg {oracle_avg:.3f} should be > 0.60"
+    assert oracle_avg > partial_avg >= floor_avg, \
+        "Score ordering violated: oracle > partial >= floor"
+    assert floor_avg == 0.001, f"Floor avg {floor_avg:.3f} should be 0.0"
     print(f"\n  ✅ Task 2: oracle({oracle_avg:.3f}) > partial({partial_avg:.3f})"
           f" ≥ random({random_avg:.3f}) ≥ floor(0.0)")
 

server/tasks/task1/environment.py

@@ -61,8 +61,6 @@ class Task1Environment(BaseEnv):
         self._done: bool = False
         self._query_history: List[str] = []
         self._seen_queries: Set[str] = set()
-        self._cost_free_steps: int = 0
-        self._decay: float = 0.0
 
     # ------------------------------------------------------------------
     # OpenEnv interface

server/tasks/task2/actions.py

@@ -19,13 +19,13 @@ PropertyRetrieverInstance = PropertyRetriever()  # Load once at module level
 def get_function_code(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle GET_FUNCTION_CODE action."""
     if ctx._is_repeated(qkey):
-        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
+        return "Repeated query.", Reward(value=ActionType.REPEATED.cost, reason="Repeated query")
 
     fn = ctx._target_fn
     code = fn.get("code", "// no code available")
     return (
         code,
-        Reward(value=-0.06, reason="get_function_code cost"),
+        Reward(value=ActionType.GET_FUNCTION_CODE.cost, reason="get_function_code cost"),
     )
 
 # TODO: Can separate comment and output_property(output_comment)
@@ -33,7 +33,7 @@ def get_function_code(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
 def get_function_natspec(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle GET_FUNCTION_NATSPEC action."""
     if ctx._is_repeated(qkey):
-        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
+        return "Repeated query.", Reward(value=ActionType.REPEATED.cost, reason="Repeated query")
 
     fn = ctx._target_fn
     name = fn["name"]
@@ -42,26 +42,26 @@ def get_function_natspec(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward
     result = f"NatSpec for '{name}':\n{natspec}"
     if out_prop:
         result += f"\n\nExpected output: {out_prop}"
-    return result, Reward(value=-0.08, reason="get_function_natspec cost")
+    return result, Reward(value=ActionType.GET_FILE_NATSPEC.cost, reason="get_function_natspec cost")
 
 
 def get_file_natspec(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle GET_FILE_NATSPEC action."""
     if ctx._is_repeated(qkey):
-        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
+        return "Repeated query.", Reward(value=ActionType.REPEATED.cost, reason="Repeated query")
 
     meta = ctx._contract.get("metadata", {})
     natspec = meta.get("natspec") or meta.get("description", "No file NatSpec available.")
     return (
         f"File NatSpec for {ctx._contract['contract_name']}:\n{natspec}",
-        Reward(value=-0.03, reason="get_file_natspec cost"),
+        Reward(value=ActionType.GET_FILE_NATSPEC.cost, reason="get_file_natspec cost"),
     )
 
 
 def get_related_functions_action(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle GET_RELATED_FUNCTIONS action."""
     if ctx._is_repeated(qkey):
-        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
+        return "Repeated query.", Reward(value=ActionType.REPEATED.cost, reason="Repeated query")
 
     name = ctx._target_fn["name"]
     related = get_related_functions(ctx._contract, name)
@@ -76,23 +76,23 @@ def get_related_functions_action(ctx: Any, qkey: str, params: Dict) -> Tuple[str
                 comment = rfn.get("comment", "")
                 summaries.append(f"  • {sig} — {comment}")
         text = f"Related functions for '{name}':\n" + "\n".join(summaries)
-    return text, Reward(value=-0.06, reason="get_related_functions cost")
+    return text, Reward(value=ActionType.GET_RELATED_FUNCTIONS.cost, reason="get_related_functions cost")
 
 
 def get_signature(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle GET_SIGNATURE action."""
     if ctx._is_repeated(qkey):
-        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
+        return "Repeated query.", Reward(value=ActionType.REPEATED.cost, reason="Repeated query")
 
     fn = ctx._target_fn
     sig = fn.get("signature")
-    return sig, Reward(value=-0.04, reason="get_signature cost")
+    return sig, Reward(value=ActionType.GET_SIGNATURE.cost, reason="get_signature cost")
 
 
 def get_similar_rule_action(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle GET_SIMILAR_RULE action."""
     if ctx._is_repeated(qkey):
-        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
+        return "Repeated query.", Reward(value=ActionType.REPEATED.cost, reason="Repeated query")
     
     PropertyRetrieverInstance.load_model()  # Ensure model is loaded before querying
     similar_rule = PropertyRetrieverInstance.get_similar_property(ctx._target_fn["code"])
@@ -101,8 +101,7 @@ def get_similar_rule_action(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Rew
             "No similar rule available for this function.",
             Reward(value=-0.20, reason="get_similar_rule cost (not found)"),
         )
-    return similar_rule, Reward(value=-0.20, reason="get_similar_rule cost")
-
+    return similar_rule, Reward(value=ActionType.GET_SIMILAR_RULE.cost, reason="get_similar_rule cost")
 
 def submit_property(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     """Handle SUBMIT_PROPERTY action for Task 2.
@@ -111,46 +110,19 @@ def submit_property(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     ---------------
     property : str  – natural-language property describing the function's behaviour
     """
-    if ctx._submitted:
-        return (
-            "❌ You have already submitted for this episode. "
-            "Only ONE submission is allowed.",
-            Reward(value=0.0, reason="Second submit_property attempt", partial=False),
-        )
- 
+    
     submitted_property = params.get("property", "").strip()
- 
     if not submitted_property:
         return (
             "submit_property requires a non-empty 'property' string in params.",
-            Reward(value=0.0, reason="Malformed submission", partial=False),
+            Reward(value=ActionType.RESUBMIT.cost, reason="Malformed submission", partial=False),
         )
  
-    ctx._submitted = True
     ctx._done      = True
+    score, confidence = ctx._grader.grade(submitted_property, ctx._step_count, ctx._cum_reward)
  
-    # grade() returns (float score in [0,1], confidence str)
-    score, confidence = ctx._grader.grade(submitted_property)  # score already in [0.0, 1.0]
-    reward_val        = float(score)                           # reward == grade for Task 2
- 
-    if confidence == "strong":
-        msg = (
-            f"✅ STRONG MATCH. Your property closely matches the target. "
-            f"Score: {score:.3f} → Reward: {reward_val:.3f}"
-        )
-    elif confidence == "moderate":
-        msg = (
-            f"🟡 MODERATE MATCH. Your property partially captures the target behaviour. "
-            f"Score: {score:.3f} → Reward: {reward_val:.3f}"
-        )
-    else:
-        msg = (
-            f"❌ LOW MATCH. Your property does not sufficiently match the target. "
-            f"Score: {score:.3f} → Reward: {reward_val:.3f}"
-        )
- 
-    return msg, Reward(
-        value=reward_val,
+    return "", Reward(
+        value=score,
         reason=f"submit_property confidence={confidence} score={score:.3f}",
         partial=False,
     )
@@ -158,7 +130,10 @@ def submit_property(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
 
 def unknown_action(ctx: Any, qkey: str, params: Dict, action_type: str) -> Tuple[str, Reward]:
     """Fallback for unknown actions."""
+    
+    ctx._done = True
     return (
-        f"Unknown action type: '{action_type}'. Valid: {[a.value for a in ActionType]}",
-        Reward(value=-0.10, reason="Unknown action"),
+        f"Unknown action type: '{action_type}'. Valid: {[a.value for a in ActionType]}, \
+        Reset environment to start again.",
+        Reward(value=ActionType.UNKNOWN.cost, reason="Unknown action"),
     )

server/tasks/task2/environment.py

@@ -27,6 +27,7 @@ from __future__ import annotations
 
 import random
 from typing import Any, Dict, List, Optional, Set
+from math import log2, floor
 
 from data.data_loader import load_contracts, sample_property_episode
 from env.base_env import BaseEnv
@@ -43,7 +44,6 @@ from .grader import Task2Grader
 from server.tasks.task2 import actions
 
 TASK_ID    = "task2_property_discovery"
-MAX_STEPS  = 15
 
 AVAILABLE_ACTIONS = [
     ActionType.GET_FUNCTION_CODE,
@@ -62,6 +62,7 @@ class Task2Environment(BaseEnv):
     def __init__(self, contracts_path: Optional[str] = None) -> None:
         self._contracts = load_contracts(contracts_path) if contracts_path else load_contracts()
         self._rng = random.Random()
+        self._max_steps: int = 40
 
         # Episode state – initialised by reset()
         self._contract:    Dict[str, Any] = {}
@@ -70,7 +71,6 @@ class Task2Environment(BaseEnv):
         self._step_count:  int = 0
         self._cum_reward:  float = 0.0
         self._done:        bool = False
-        self._submitted:   bool = False          # only one submit_property allowed
         self._query_hist:  List[str] = []
         self._seen:        Set[str] = set()
 
@@ -86,6 +86,7 @@ class Task2Environment(BaseEnv):
         self._grader = Task2Grader(
             function_name=self._target_fn["name"],
             property=self._target_fn["property"],
+            n = floor(log2(len(self._contract["functions"])))
         )
         self._step_count = 0
         self._cum_reward = 0.0
@@ -110,6 +111,9 @@ class Task2Environment(BaseEnv):
     def step(self, action: Action) -> StepResult:
         if self._done:
             raise RuntimeError("Episode is done. Call reset() to start a new episode.")
+        
+        if self._step_count > self._max_steps:
+            raise RuntimeError("Exceeded maximum number of steps allowed. Call reset() to start a new episode.")
 
         self._step_count += 1
         result_text, reward = self._dispatch(action)
@@ -147,12 +151,8 @@ class Task2Environment(BaseEnv):
         return Observation(
             task_id=TASK_ID,
             contract_name=self._contract.get("contract_name", ""),
-            contract_description=self._contract.get("metadata", {}).get("description", ""),
-            available_actions=[a.value for a in AVAILABLE_ACTIONS],
             last_action=last_action,
             last_action_result=last_result,
-            step_count=self._step_count,
-            cumulative_reward=self._cum_reward,
             done=self._done,
             extra={
                 "target_function": self._target_fn.get("name", ""),
@@ -181,8 +181,6 @@ class Task2Environment(BaseEnv):
         params = action.params
         qkey = self._qkey(at, params)
 
-        # Mapping from ActionType to handler function
-        # Each handler expects (ctx, qkey, params) and returns (str, Reward)
         handlers = {
             ActionType.GET_FUNCTION_CODE:       actions.get_function_code,
             ActionType.GET_FUNCTION_NATSPEC:    actions.get_function_natspec,

server/tasks/task2/grader.py

@@ -9,12 +9,6 @@ Grade range: 0.0 – 1.0  (matchscore output, already normalised).
 from typing import Tuple
 from utils import SemanticMatcher
 
-_SCORE_MIN = 0.001   # grades are strictly (0, 1)
-_SCORE_MAX = 0.999
-
-def _clamp(v: float) -> float:
-    return max(_SCORE_MIN, min(_SCORE_MAX, v))
-
 class Task2Grader:
     """
     Grades a Task 2 property submission.
@@ -25,15 +19,20 @@ class Task2Grader:
     property      : the 'property' field from the target function's data
     """
 
-    def __init__(self, function_name: str, property: str) -> None:
+    def __init__(self, function_name: str, property: str, n: int) -> None:
         self.function_name = function_name
         self.property      = property
+        self.n             = n
+        self._decay        = 0.75
 
-    def grade(self, submitted: str) -> Tuple[float, str]:
+    def grade(self, submitted: str, steps: int, cummulative_cost: int) -> Tuple[float, str]:
         """Deterministic grade strictly in (0, 1)."""
         if not submitted or not submitted.strip():
-            return _clamp(0.0), "no_match"   # → 0.001
+            return 0.001, "no_match"
 
         matcher = SemanticMatcher()
-        score   = matcher.matchscore(self.property, submitted)  # already clamped by SemanticMatcher
-        return _clamp(score), matcher.confidence()
+        match_score   = matcher.matchscore(self.property, submitted)
+        free_budget = (cummulative_cost / steps) * (self.n + 2)
+        final_score = match_score * (self._decay ** max(0, cummulative_cost - free_budget))
+        
+        return final_score, matcher.confidence()