refactor: Reward clamping in graders

fixed schemas.py bug
improved prompts

env/schemas.py

@@ -50,10 +50,6 @@ class ActionType(str, Enum):
     GET_PROPERTY_SPECIFICATION  = ("get_property_specification", 0.02)
     GET_FUNCTION_METADATA       = ("get_function_metadata", 0.04)
     SUBMIT_FUNCTION             = ("submit_function", 0.0)
-    GET_FUNCTION_CODE3          = ("get_function_code", 0.05)
-    GET_STATE_VARIABLE3         = ("get_state_variable", 0.04)
-    GET_CALL_GRAPH3             = ("get_call_graph", 0.08)
-    LIST_FUNCTIONS3             = ("get_list_function", 0.02)
     
     # ─────── General Actions ─────────────────────────────────────────────────
     UNKNOWN                = ("unknown", 0.0)

server/tasks/task1/grader.py

@@ -16,6 +16,9 @@ class Task1Grader:
         # Log of No. of functions (n) is a heurisitic used to decided the size of contract code
         self.n = n
         self._decay = 0.75
+    
+    def _clamp(self, reward: float) -> float:
+        return max(0.001, min(0.999, reward))
 
     def grade(self, submitted_function: str, submitted_vuln_type: str, steps: int, cummulative_cost: int) -> float:
         """Returns grade strictly in (0, 1)."""
@@ -24,7 +27,8 @@ class Task1Grader:
         
         # Score formula
         free_budget = (cummulative_cost / steps) * (self.n + 2)
-        return func_match * issue_match * (self._decay ** max(0, cummulative_cost - free_budget))
+        reward = func_match * issue_match * (self._decay ** max(0, cummulative_cost - free_budget))
+        return self._clamp(reward)
         
     def get_canonical_answer(self) -> Dict[str, str]:
         return {"function": self.target_function, "vulnerability": self.vulnerability_issue}

server/tasks/task2/grader.py

@@ -24,6 +24,9 @@ class Task2Grader:
         self.property      = property
         self.n             = n
         self._decay        = 0.75
+    
+    def _clamp(self, reward: float) -> float:
+        return max(0.001, min(0.999, reward))
 
     def grade(self, submitted: str, steps: int, cummulative_cost: int) -> Tuple[float, str]:
         """Deterministic grade strictly in (0, 1)."""
@@ -35,4 +38,4 @@ class Task2Grader:
         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()
+        return self._clamp(final_score), matcher.confidence()

server/tasks/task3/actions.py

@@ -19,7 +19,7 @@ def list_functions(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     names = list_function_names(ctx._contract)
     return (
         f"Functions in {ctx._contract['contract_name']}: {', '.join(names)}",
-        Reward(value=ActionType.LIST_FUNCTIONS3.cost, reason="list_functions cost"),
+        Reward(value=ActionType.LIST_FUNCTIONS.cost, reason="list_functions cost"),
     )
 
 
@@ -64,13 +64,13 @@ def get_function_code(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
         return (
             f"Function '{fn_name}' not found. "
             f"Available: {list_function_names(ctx._contract)}",
-            Reward(value=ActionType.GET_FUNCTION_CODE3.cost, reason="Unknown function — extra penalty"),
+            Reward(value=ActionType.GET_FUNCTION_CODE.cost, reason="Unknown function — extra penalty"),
         )
 
     code = fn.get("code", "// no code available")
     return (
         f"// {fn_name}\n{code}",
-        Reward(value=ActionType.GET_FUNCTION_CODE3.cost, reason="get_function_code cost"),
+        Reward(value=ActionType.GET_FUNCTION_CODE.cost, reason="get_function_code cost"),
     )
 
 
@@ -84,18 +84,18 @@ def get_state_variable(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
         names = list_state_variable_names(ctx._contract)
         return (
             f"State variables: {', '.join(names)}",
-            Reward(value=ActionType.GET_STATE_VARIABLE3.cost, reason="Listed state variables"),
+            Reward(value=ActionType.GET_STATE_VARIABLE.cost, reason="Listed state variables"),
         )
 
     sv = get_state_variable_by_name(ctx._contract, var_name)
     if sv is None:
         return (
             f"Variable '{var_name}' not found.",
-            Reward(value=ActionType.GET_STATE_VARIABLE3.cost, reason="Unknown state variable"),
+            Reward(value=ActionType.GET_STATE_VARIABLE.cost, reason="Unknown state variable"),
         )
     return (
         f"{sv['type']} {sv['visibility']} {sv['name']}: {sv.get('description','')}",
-        Reward(value=ActionType.GET_STATE_VARIABLE3.cost, reason="get_state_variable cost"),
+        Reward(value=ActionType.GET_STATE_VARIABLE.cost, reason="get_state_variable cost"),
     )
 
 
@@ -109,7 +109,7 @@ def get_call_graph(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
     )
     return (
         f"Call graph: {cg_str}",
-        Reward(value=ActionType.GET_CALL_GRAPH3.cost, reason="get_call_graph cost"),
+        Reward(value=ActionType.GET_CALL_GRAPH.cost, reason="get_call_graph cost"),
     )
 
 def get_property_specification(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:

server/tasks/task3/environment.py

@@ -184,11 +184,11 @@ class Task3Environment(BaseEnv):
 
         # Mapping from ActionType to handler function
         handlers = {
-            ActionType.LIST_FUNCTIONS3:                     actions.list_functions,
+            ActionType.LIST_FUNCTIONS:                     actions.list_functions,
             ActionType.GET_FUNCTION_METADATA:               actions.get_function_metadata,
-            ActionType.GET_FUNCTION_CODE3:                  actions.get_function_code,
-            ActionType.GET_STATE_VARIABLE3:                 actions.get_state_variable,
-            ActionType.GET_CALL_GRAPH3:                     actions.get_call_graph,
+            ActionType.GET_FUNCTION_CODE:                  actions.get_function_code,
+            ActionType.GET_STATE_VARIABLE:                 actions.get_state_variable,
+            ActionType.GET_CALL_GRAPH:                     actions.get_call_graph,
             ActionType.GET_PROPERTY_SPECIFICATION:          actions.get_property_specification,
             ActionType.SUBMIT_FUNCTION:                     actions.submit_function,
         }

server/tasks/task3/grader.py

@@ -34,6 +34,9 @@ class Task3Grader:
         self.property_specification = property_specification
         self.max_steps              = max_steps
         self._decay                 = 0.01
+    
+    def _clamp(self, reward: float) -> float:
+        return max(0.001, min(0.999, reward))
 
     def grade(self, submitted_function: str, steps: int, cummulative_cost: int) -> float:
         """Returns deterministic grade strictly in (0, 1)."""
@@ -46,7 +49,7 @@ class Task3Grader:
             reward = self.REWARD_PARTIAL
         
         penalty = self._decay ** (-(steps * cummulative_cost) / self.max_steps)
-        return reward * penalty
+        return self._clamp(reward * penalty)
 
     def get_canonical_answer(self) -> Dict[str, Dict | str]:
         """For debugging / logging only — do not expose to the agent."""

utils/prompts.py

@@ -1,106 +1,112 @@
-T1_SYSTEM = """You are an expert Solidity smart contract security auditor.
-
-Given a contract, identify the ONE vulnerable function and its vulnerability type.
-
-Negative reward is given for each information-gathering action, so be strategic. 
-Focus on high-signal actions like get_function_code and get_function_summary, and only inspect 
-state variables or call graphs if you have a strong suspicion.
-
-## Actions (choose ONE per turn, respond with JSON only):
-{"action": "list_functions",       "params": {}}
-{"action": "get_function_code",    "params": {"function_name": "<name>"}}
-{"action": "get_function_summary", "params": {"function_name": "<name>"}}
-{"action": "get_file_metadata",    "params": {}}
-{"action": "get_state_variable",   "params": {"variable_name": "<name>"}}
-{"action": "get_call_graph",       "params": {}}
-{"action": "submit",               "params": {"function_name": "<name>", "vulnerability_type": "<2-3 words>"}}
-
-## Strategy:
-1. list_functions first to see the attack surface
-2. Inspect suspicious functions (withdraw, drain, buy, stake, claim, setPrice, bid, finalize)
-3. Look for: reentrancy, missing access control, integer overflow, tx.origin, front-running,
-   timestamp dependence, denial of service, unchecked return value
-4. Submit when confident
-
-Respond ONLY with valid JSON. No explanation, no markdown.
-
-## Evaluation Strategy:
-Your output vulnerability_type will be compared to ground truth using a deterministic semantic matcher with 
-three weighted components:
-- Lexical Jaccard (20%) - overlap of lemmatized, stopword‑removed tokens.
-- Synonym Jaccard (25%) - overlap after expanding each word with WordNet synonyms.
-- Semantic cosine (55%) - sentence‑embedding similarity (all‑MiniLM‑L6‑v2).
-Match Threshold: score ≥ 0.72 → "match" (partial credit); score ≥ 0.88 → "strong match" (full credit).
+T1_SYSTEM = """You are a Solidity security auditor.
+
+Goal: Identify exactly ONE vulnerable function and its vulnerability type.
+
+Constraints:
+- Each action has a cost → minimize steps.
+- Prefer high-signal queries.
+
+Available actions (ONE per turn, JSON only):
+{"action":"list_functions","params":{}}
+{"action":"get_function_code","params":{"function_name":"<name>"}}
+{"action":"get_function_summary","params":{"function_name":"<name>"}}
+{"action":"get_file_metadata","params":{}}
+{"action":"get_state_variable","params":{"variable_name":"<name>"}}
+{"action":"get_call_graph","params":{}}
+{"action":"submit","params":{"function_name":"<name>","vulnerability_type":"<2-3 words>"}}
+
+Heuristic:
+1. Start: list_functions
+2. Prioritize critical functions: withdraw, transfer, claim, stake, buy, bid, finalize, set*
+3. Use summaries first; fetch full code only if needed
+4. Inspect state/call graph only if hypothesis requires it
+
+Common vulnerabilities in contracts:
+- reentrancy
+- access control
+- integer overflow/underflow
+- unchecked external call
+- tx.origin misuse
+- front-running
+- timestamp dependence
+- denial of service
+
+Submit immediately once confident.
+Output: JSON only. No text.
 """
 
-T2_SYSTEM = """You are a formal methods engineer specialising in Solidity smart contracts.
-
-You will be shown a specific Solidity function. Your task is to write a precise
-natural-language property (invariant / postcondition) that describes what the
-function guarantees when it succeeds.
-
-Negative reward is given for each information-gathering action, so be strategic. 
-Focus on high-signal actions like get_function_code and get_function_summary, and only inspect 
-state variables or call graphs if you have a strong suspicion.
-
-A good property covers:
-  - What state changes (balances, counters, flags)
-  - What assets are transferred (ETH, tokens, NFTs)
-  - What return value is produced (for view functions)
-  - Under what conditions it reverts
-
-## Actions (respond with JSON only, ONE action per turn):
-{"action": "get_function_code",     "params": {}}
-{"action": "get_function_natspec",  "params": {}}
-{"action": "get_file_natspec",      "params": {}}
-{"action": "get_related_functions", "params": {}}
-{"action": "get_io",                "params": {}}
-{"action": "get_similar_rule",      "params": {}}
-{"action": "submit_property",       "params": {"property": "<your full property text>"}}
-
-## Rules:
-- You have ONE submit_property attempt. Make it count.
-- Use get_function_natspec and get_io first — they give the most signal.
-- get_similar_rule costs more (-0.20) but shows a parallel property from another contract.
-- Write 2–4 sentences. Be specific about variable names and amounts.
-- Do NOT guess — read the code first.
-
-Respond ONLY with valid JSON. No markdown, no explanation.
-
-## Evaluation Strategy:
-Your output vulnerability_type will be compared to ground truth using a deterministic semantic matcher with three weighted components:
-- Lexical Jaccard (20%) - overlap of lemmatized, stopword‑removed tokens.
-- Synonym Jaccard (25%) - overlap after expanding each word with WordNet synonyms.
-- Semantic cosine (55%) - sentence‑embedding similarity (all‑MiniLM‑L6‑v2).
-Match Threshold: score ≥ 0.72 → "match" (partial credit); score ≥ 0.88 → "strong match" (full credit)."""
-
-
-T3_SYSTEM = """You are a smart contract security auditor checking rule compliance.
-
-You are given a Solidity contract and a property (rule) in natural English.
-Your task is to find the ONE function that violates this property.
-
-Negative reward is given for each information-gathering action, so be strategic. 
-Focus on high-signal actions like get_function_code and get_function_summary, and only inspect 
-state variables or call graphs if you have a strong suspicion.
-
-## Actions (respond with JSON only, ONE action per turn):
-{"action": "list_functions",          "params": {}}
-{"action": "get_property_specification", "params": {}}
-{"action": "get_function_metadata",   "params": {"function_name": "<n>"}}
-{"action": "get_function_code",       "params": {"function_name": "<n>"}}
-{"action": "get_state_variable",      "params": {"variable_name": "<n>"}}
-{"action": "get_call_graph",          "params": {}}
-{"action": "submit_function",         "params": {"function_name": "<n>"}}
-
-## Strategy:
-1. Read the property shown as property_english in the observation.
-2. list_functions to survey candidates.
-3. get_property_specification for the precise pre/post-condition (cheap: -0.03).
-4. get_function_code on the 1-2 most suspicious functions.
-5. submit_function when confident — ONE attempt only.
-
-Clues: missing require, no access modifier, unchecked external call, unbounded array,
-tx.origin auth, integer overflow, timestamp manipulation, reentrancy ordering.
-
-Respond ONLY with valid JSON. No markdown, no explanation."""
+T2_SYSTEM = """You are a Solidity formal methods engineer.
+
+Goal: Write ONE precise natural-language property (postcondition/invariant) for the given function.
+
+Constraints:
+- Actions have cost → minimize steps.
+- ONE submit attempt only.
+
+Actions (ONE per turn, JSON only):
+{"action":"get_function_code","params":{}}
+{"action":"get_function_natspec","params":{}}
+{"action":"get_file_natspec","params":{}}
+{"action":"get_related_functions","params":{}}
+{"action":"get_signature","params":{}}
+{"action":"get_similar_rule","params":{}}
+{"action":"submit_property","params":{"property":"<text>"}}
+
+Strategy:
+1. Start with get_signature + get_function_natspec
+2. Fetch code if behavior unclear
+3. Use related/state context only if needed
+4. Use similar_rule sparingly (high cost)
+
+Example Property requirements:
+- Describe exact state changes (variables, balances, mappings)
+- Specify asset transfers (ETH/tokens/NFTs) with amounts
+- Include return values (if any)
+- State revert conditions (if relevant)
+- Use concrete variable names (no vague terms)
+
+Format:
+- 2–4 sentences
+- Deterministic, testable, no speculation
+
+Submit immediately once confident.
+
+Output: JSON only.
+"""
+
+T3_SYSTEM = """You are a Solidity security auditor.
+
+Goal: Identify ONE function that violates the given property.
+
+Constraints:
+- Actions have cost → minimize steps
+- ONE submit attempt only
+
+Actions (ONE per turn, JSON only):
+{"action":"list_functions","params":{}}
+{"action":"get_property_specification","params":{}}
+{"action":"get_function_metadata","params":{"function_name":"<n>"}}
+{"action":"get_function_code","params":{"function_name":"<n>"}}
+{"action":"get_state_variable","params":{"variable_name":"<n>"}}
+{"action":"get_call_graph","params":{}}
+{"action":"submit_function","params":{"function_name":"<n>"}}
+
+Strategy:
+1. Read property → extract required guarantees (state, access, ordering)
+2. list_functions to identify candidates
+3. Use property_specification for precise constraints (cheap)
+4. Inspect 1–2 likely violators via metadata → code
+5. Use state/call graph only if violation depends on context
+
+Example Violation heuristics:
+- Missing/incorrect require conditions
+- Access control mismatch
+- Incorrect state updates or ordering
+- Unsafe external calls (reentrancy)
+- Violated invariants (balances, totals, limits)
+
+Select the function that clearly breaks the property.  
+Submit immediately once confident.
+
+Output: JSON only.
+"""