SmartContractAudit / README.md
ajaxwin
refactor: Update task configurations and grading logic for improved scoring and consistency
dccaaac
metadata
title: Smart Contract Audit RL Environment
emoji: πŸ”
colorFrom: blue
colorTo: indigo
sdk: docker
app_port: 7860
tags:
  - openenv
  - reinforcement-learning
  - smart-contracts
  - solidity
  - security
  - evaluation
  - openenv
license: mit
short_description: OpenEnv RL environment for smart contract security auditing

πŸ” Smart Contract Audit RL Environment

An OpenEnv-compliant reinforcement learning environment for training and evaluating AI agents on real-world Solidity smart contract security auditing tasks.


Overview

Smart contract auditing is a high-stakes, expert-level task performed by professional security researchers. Mistakes cost millions β€” the Ethereum ecosystem has lost over $3 billion to exploits in audited and unaudited contracts alike. This environment simulates the core reasoning loop of a smart contract auditor, enabling RL agents to learn structured exploration strategies for vulnerability detection, property discovery, and rule checking.

The dataset is derived from real audit reports published by Certora, covering three production-grade DeFi protocols:

Source Protocol
Certora Audit AaveVault
Certora Audit AaveVaultV2
Certora Audit Lido Finance

Each episode exposes a fragment of a real Solidity contract. The agent must use a structured action API β€” mirroring how a human auditor would methodically inspect a codebase β€” to accomplish a defined objective within a fixed step budget.


Environment Architecture

SmartContractEnv/
β”œβ”€β”€ agents/
β”‚   β”œβ”€β”€ task1.py
β”‚   β”œβ”€β”€ task2.py
β”‚   └── task3.py
β”œβ”€β”€ data/
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ contracts.json
β”‚   β”œβ”€β”€ data_loader.py
β”‚   β”œβ”€β”€ properties.csv
β”‚   β”œβ”€β”€ Template.json
β”‚   β”œβ”€β”€ vulnerabilities.json
β”‚   └── vulnerabilities.md
β”œβ”€β”€ env/
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ base_env.py
β”‚   └── schemas.py
β”œβ”€β”€ server/
β”‚   β”œβ”€β”€ tasks/
β”‚   β”‚   β”œβ”€β”€ task1/
β”‚   β”‚   β”‚   β”œβ”€β”€ __init__.py
β”‚   β”‚   β”‚   β”œβ”€β”€ actions.py
β”‚   β”‚   β”‚   β”œβ”€β”€ environment.py
β”‚   β”‚   β”‚   └── grader.py
β”‚   β”‚   β”œβ”€β”€ task2/
β”‚   β”‚   β”‚   β”œβ”€β”€ __init__.py
β”‚   β”‚   β”‚   β”œβ”€β”€ actions.py
β”‚   β”‚   β”‚   β”œβ”€β”€ environment.py
β”‚   β”‚   β”‚   └── grader.py
β”‚   β”‚   β”œβ”€β”€ task3/
β”‚   β”‚   β”‚   β”œβ”€β”€ __init__.py
β”‚   β”‚   β”‚   β”œβ”€β”€ actions.py
β”‚   β”‚   β”‚   β”œβ”€β”€ environment.py
β”‚   β”‚   β”‚   └── grader.py
β”‚   β”‚   └── __init__.py
β”‚   β”œβ”€β”€ __init__.py
β”‚   └── app.py
β”œβ”€β”€ utils/
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ prompts.py
β”‚   β”œβ”€β”€ propertyretriever.py
β”‚   └── semanticmatcher.py
β”œβ”€β”€ .env
β”œβ”€β”€ .gitignore
β”œβ”€β”€ demo.py
β”œβ”€β”€ Dockerfile
β”œβ”€β”€ eval.py
β”œβ”€β”€ inference.py
β”œβ”€β”€ LICENSE.txt
β”œβ”€β”€ openenv.yaml
β”œβ”€β”€ pyproject.toml
β”œβ”€β”€ README.md
β”œβ”€β”€ requirements.txt
└── validate-submission.sh

Tasks

Task 1 β€” Targeted Vulnerability Detection (Medium)

Real-world analogue: A security auditor is handed a Solidity file and asked to pinpoint the vulnerable function and describe the class of bug.

Setup: The agent receives a single Solidity file. The episode selects one vulnerable function at random from the dataset (7–8 available) on each reset().

Objective: Identify the vulnerable function and describe its issue in 2–3 words (e.g., "reentrancy", "integer overflow", "unchecked return value"). Submit "NO" if no vulnerability exists.

Action Space:

Action Reward Notes
list_functions βˆ’0.05 Returns all function signatures in the file
get_function_code βˆ’0.10 (wrong fn) / +0.05 (correct fn) Returns raw Solidity source of one function
get_function_summary βˆ’0.05 (wrong) / +0.03 (correct) Returns NatSpec comments for a function
get_file_metadata βˆ’0.04 Returns the file's header comment / pragma / imports
get_state_variables βˆ’0.05 Returns all contract-level state variable declarations
get_call_graph βˆ’0.08 Returns the inter-function call graph
get_task_state 0.00 Returns current step count and cumulative reward
submit +5.00 (correct) / βˆ’1.50 (wrong) One submission allowed per episode
(repeated query) βˆ’0.40 Penalty for querying the exact same action+params twice
(unknown action) βˆ’0.20 Any unrecognised action type

Episode terminates on submit or when the step budget is exhausted.


Task 2 β€” Property Discovery (Hard)

Real-world analogue: A formal verification engineer must derive an invariant or safety property for a contract function β€” the kind written as a Certora Verification Language (CVL) spec.

Setup: The agent receives a single function extracted from a Solidity file, along with a brief description of the broader contract. The episode targets a function that has a known, labelled property in the dataset.

Objective: Produce a natural-language description of the function's key safety property (e.g., "The total shares minted must never exceed the total underlying assets deposited").

Action Space:

Action Reward Notes
get_file_natspec βˆ’0.03 File-level NatSpec documentation
get_function_natspec βˆ’0.08 Function-level NatSpec comments
get_function_code βˆ’0.06 Raw Solidity source of the target function
get_related_functions βˆ’0.06 Functions that call or are called by the target
get_input_output βˆ’0.04 Parameter names/types and return values
get_similar_property βˆ’0.20 Hard-coded reference property from a different contract
submit_property 0–5 (graded) One attempt per episode. Scored by deterministic similarity checker

Grading: Submission reward is computed by a deterministic checker that combines keyword overlap and structural similarity against the ground-truth property. Score is normalised to [0, 5] and then scaled to [0.0, 1.0] for the episode return.


Task 3 β€” Rule Checker (Easy)

Real-world analogue: Given a known security rule (e.g., "functions that transfer funds must emit a Transfer event"), identify which function in the contract violates it.

Setup: The agent receives a Solidity file and a natural-language description of a property/rule. At least one function in the file violates this rule.

Objective: Identify the name of the rule-breaking function.

Action Space:

Action Reward Notes
get_property_specification βˆ’0.03 Returns a pseudo-formal (CVL-like) version of the property
list_functions βˆ’0.05 All function signatures in the file
get_function_metadata βˆ’0.05 Visibility, modifiers, and signature for a function
get_function_code βˆ’0.10 Raw Solidity source of one function
get_state_variables βˆ’0.05 Contract-level state variable declarations
get_call_graph βˆ’0.08 Inter-function call graph
submit +5.00 (exact) / +1.50 (sub-caller) / βˆ’1.50 (wrong) One submission per episode

Partial credit: If the agent names a function that calls the true violating function, it receives +1.50 rather than the full +5.00. This rewards reasoning that reaches the right area of the call graph.


Reward Design

Rewards are shaped to encourage efficient, targeted exploration and discourage two failure modes: aimless browsing and brute-force guessing.

R_episode = Ξ£(step_rewards) + final_submission_reward
  • Exploration costs are small and graduated by information value β€” cheap actions (metadata) cost less than expensive ones (full code retrieval).
  • Correct-direction bonuses on get_function_code in Task 1 reward navigating toward the vulnerable function before committing.
  • Repetition penalty (βˆ’0.40) discourages looping over the same queries.
  • Wrong submission (βˆ’1.50) is painful enough to deter random guessing but recoverable through efficient prior exploration.
  • Episode score is normalised to [0.0, 1.0] for the OpenEnv grader: score = max(0, R_episode) / 5.0.

Observation Space

Every step() and reset() returns a typed Observation object:

class Observation(BaseModel):
    task_id: str                        # "task1_vuln_detection" | "task2_property_discovery" | "task3_rule_checker"
    step: int                           # Current step index (0-indexed)
    max_steps: int                      # Episode step budget
    cumulative_reward: float            # Running reward total
    done: bool                          # Episode terminal flag
    content: str                        # Main textual payload (code, summary, error, etc.)
    metadata: dict[str, Any]            # Extra context (function name, contract name, etc.)
    initial_description: str            # Persistent contract/task description shown every step

Action Space

Actions are typed Action objects passed to step():

class Action(BaseModel):
    action_type: str                    # One of the action names listed per task above
    params: dict[str, str]             # e.g. {"function_name": "withdraw"}

All unknown action_type values return a penalty observation without terminating the episode.


OpenEnv Interface

The environment exposes a standard HTTP API:

Method Path Description
GET /health Liveness probe β€” returns {"status": "ok"}
GET /tasks Lists all tasks with ID, difficulty, and status
POST /reset Starts a new episode. Body: {"task_id": str, "seed": int}
POST /step Takes one action. Body: {"action_type": str, "params": {}}
GET /state Returns full internal episode state (debug)
GET /action_space Returns JSON schema of valid actions
GET /observation_space Returns JSON schema of observation structure

Quick Start

SPACE_URL=http://localhost:7860

# Start a new episode for Task 1
curl -X POST $SPACE_URL/reset \
  -H "Content-Type: application/json" \
  -d '{"task_id": "task1_vuln_detection", "seed": 42}'

# List all functions in the contract
curl -X POST $SPACE_URL/step \
  -H "Content-Type: application/json" \
  -d '{"action_type": "list_functions", "params": {}}'

# Inspect a specific function
curl -X POST $SPACE_URL/step \
  -H "Content-Type: application/json" \
  -d '{"action_type": "get_function_code", "params": {"function_name": "withdraw"}}'

# Submit your answer
curl -X POST $SPACE_URL/step \
  -H "Content-Type: application/json" \
  -d '{"action_type": "submit", "params": {"function_name": "withdraw", "vulnerability_type": "reentrancy"}}'

Setup & Installation

Prerequisites

  • Docker β‰₯ 20.10
  • Python 3.11+ (for local development)
  • OPENAI_API_KEY, API_BASE_URL, MODEL_NAME, HF_TOKEN environment variables

Run with Docker

# Build the image
docker build -t sc-audit-env .

# Run the container
docker run -p 7860:7860 \
  -e OPENAI_API_KEY=$OPENAI_API_KEY \
  -e API_BASE_URL=$API_BASE_URL \
  -e MODEL_NAME=$MODEL_NAME \
  sc-audit-env

# Verify it's running
curl http://localhost:7860/health

Run Locally (Development)

pip install -r requirements.txt
uvicorn server.app:app --host 0.0.0.0 --port 7860 --reload

Baseline Inference Script

The inference.py script runs an OpenAI-compatible model against all three tasks and reports episode scores. It reads credentials from environment variables and completes in under 20 minutes on a 2 vCPU / 8 GB machine.

export OPENAI_API_KEY=your_key
export API_BASE_URL=your custom endpoint
export MODEL_NAME=your custom model

python inference.py

Expected output:

=== Smart Contract Audit RL β€” Baseline Evaluation ===

Task 1 | Targeted Vulnerability Detection  | Score: 0.41 | Steps used: 8/15
Task 2 | Property Discovery                | Score: 0.28 | Steps used: 6/10
Task 3 | Rule Checker                      | Score: 0.72 | Steps used: 4/10

Overall average: 0.47

Note: Scores are stochastic due to random episode selection on reset(). Run with a fixed seed (--seed 42) for reproducible results.

Agent System Prompt

The inference script injects the following system prompt to guide output format:

You are a smart contract security auditor. You will be given access to a Solidity
contract via a structured action API. Use the available actions to investigate the
contract, then submit your answer.

Always respond with a single JSON object:
{"action_type": "<action>", "params": {"<key>": "<value>"}}

Do not include any other text outside the JSON object.

openenv.yaml

name: smart-contract-audit-env
version: "1.2.0"
description: >
  OpenEnv RL environment for Solidity smart contract security auditing.
  Agents explore real-world DeFi contracts using a structured action API
  to detect vulnerabilities, discover properties, and check rule compliance.
tasks:
  - id: task1_vuln_detection
    name: Targeted Vulnerability Detection
    difficulty: medium
    max_steps: 40
    max_score: 1.0
  - id: task2_property_discovery
    name: Property Discovery
    difficulty: hard
    max_steps: 40
    max_score: 1.0
  - id: task3_rule_checker
    name: Rule Checker
    difficulty: easy
    max_steps: 20
    max_score: 1.0
observation_schema: models/observation.py
action_schema: models/action.py
app_port: 7860

Data

The dataset (data/dataset.json) contains 7–8 labelled entries per contract, each with format accoding to data/template.json: Ground truth is never exposed to the agent via any action. The submit action is the only path to positive reward.


Design Notes & Known Limitations

  • Reward calibration: Step penalties and submission rewards may need tuning based on empirical agent performance. Current values are derived from initial design rationale, not from extensive ablation.
  • Call graph granularity: The current get_call_graph action returns the entire graph at once. A future revision could expose it incrementally (per-function neighbours) to make the action more informative and cost-proportional.
  • Vulnerability naming: Vulnerability types do not follow a fixed taxonomy. Grading uses keyword + semantic matching against a curated synonym list (e.g., "re-entrancy" ≑ "reentrancy").
  • Dataset size: The current dataset covers 3 contracts with 7–8 vulnerabilities each. Expanding to more Certora audit reports would improve task diversity and reduce overfitting risk.
  • get_function_code decomposition: This action could be split into finer-grained sub-actions (get_parameters, get_return_values, get_modifiers) to give agents a more gradual information ladder.
  • Property similarity scoring (Task 2): Sentence transformer models cannot be used in the containerised environment due to memory constraints. The checker instead uses TF-IDF cosine similarity combined with keyword matching against the ground-truth property.

License

MIT β€” see LICENSE for details.

Data sourced from public Certora audit reports. Solidity source files are reproduced for research and evaluation purposes.


Citation

@misc{sc-audit-openenv-2025,
  title   = {Smart Contract Audit RL Environment},
  year    = {2025},
  note    = {OpenEnv-compliant RL environment for Solidity security analysis.
             Data sourced from Certora audit reports (AaveVault, AaveVaultV2, Lido Finance).}
}