Meta OpenEnv Submission: AegisGym v1.0.0

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 8000
+
+# Start OpenEnv API server
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,56 @@
+# AegisGym: Financial Compliance & AML Sandbox
+
+AegisGym is a reinforcement learning sandbox for training and evaluating autonomous financial auditors. It simulates real-world banking compliance tasks, including Sanction Checks, Anti-Money Laundering (AML) detection, and Regulatory Alignment.
+
+## 🏦 Motivation
+Financial institutions process millions of transactions daily. Human auditors often struggle with "smurfing" hidden in noise or complex regulatory clauses across jurisdictions. AegisGym provides a rigorous environment to train LLM-based agents to detect financial crime with high explainability and logical reasoning.
+
+## 🛠️ Environment Specification
+
+### 📝 Action Space
+The agent must provide an `AuditAction`:
+- `action_type`: `APPROVE`, `FLAG`, or `REQUEST_INFO`.
+- `target_id`: The ID of the account or transaction.
+- `regulation_citation`: A string citing the relevant regulation (e.g., "BSA-31-USC-5318").
+
+### 👁️ Observation Space
+The agent receives an `AuditObservation`:
+- `transactions`: A list of recent transaction dictionaries.
+- `account_metadata`: Details about the account age, tier, and history.
+- `retrieved_regs`: Relevant regulatory guidelines fetched via RAG.
+
+### 🎯 Tasks & Difficulty
+| Task ID | Name | Difficulty | Description |
+|---------|------|------------|-------------|
+| `easy_audit` | Sanction Check | Easy | Identify accounts on a blocklist. |
+| `medium_audit` | Smurfing Detection | Medium | Detect structuring patterns under withdrawal limits. |
+| `hard_audit` | Regulatory Alignment | Hard | Accurately cite complex regulations for high-risk tx. |
+
+## 🚀 Setup & Usage
+
+### Prerequisites
+- Python 3.10+
+- `openenv-core`
+- Hugging Face Space (for deployment)
+
+### Installation
+```bash
+pip install -r requirements.txt
+```
+
+### Running Locally
+```bash
+uvicorn app:app --host 0.0.0.0 --port 7860
+```
+
+### Baseline Inference
+```bash
+export API_BASE_URL="https://api.openai.com/v1"
+export MODEL_NAME="gpt-4o"
+export OPENAI_API_KEY="your_key"
+python inference.py
+```
+
+## 🐋 Deployment
+This environment is designed for Hugging Face Spaces. Use the provided `Dockerfile`.
+- **Public URL:** [armaan020/AegisGym](https://huggingface.co/spaces/armaan020/AegisGym)

app.py

@@ -0,0 +1,5 @@
+from openenv.core.env_server import create_fastapi_app
+from server import AegisGymEnv
+from models import AuditAction, AuditObservation
+
+app = create_fastapi_app(AegisGymEnv, AuditAction, AuditObservation)

baseline_inference.py

@@ -0,0 +1,52 @@
+import os
+from server import AegisGymEnv
+from models import AuditAction
+
+def main():
+    print("--- AegisGym Final Baseline Inference Simulation (Official API) ---")
+    
+    env = AegisGymEnv()
+    obs = env.reset()
+    
+    total_score = 0.0
+    episodes = 10
+    
+    for i in range(episodes):
+        tier = env.state.current_tier
+        print(f"\nEpisode {i+1} | Tier: {tier.upper()}")
+        
+        # Mocking an agent's decision based on the tier
+        if tier == "easy":
+            action = AuditAction(
+                action_type="FLAG",
+                target_id="ACC-BL-001",
+                regulation_citation="Sanctions List Match"
+            )
+        elif tier == "medium":
+            action = AuditAction(
+                action_type="FLAG",
+                target_id="ACC-SMURF-99",
+                regulation_citation="Structuring/Smurfing 9000-10000 range"
+            )
+        else:
+            action = AuditAction(
+                action_type="FLAG",
+                target_id="ACC-REG-VIOLATOR",
+                regulation_citation="EU-AI-Act-Art-57"
+            )
+            
+        # step() returns an AuditObservation instance only
+        obs = env.step(action)
+        print(f"Action Taken: {action.action_type} for {action.target_id}")
+        print(f"Reward: {obs.reward} | Done: {obs.done}")
+        
+        total_score += obs.reward
+        if obs.done:
+            break
+            
+    print(f"\n--- Reproducibility Report ---")
+    print(f"Total Episodes: {min(episodes, env.state.step_count)}")
+    print(f"Custom Agent Mean Score (Reward): {total_score / min(episodes, env.state.step_count)}")
+
+if __name__ == "__main__":
+    main()

check_consistency.py

@@ -0,0 +1,28 @@
+import os
+
+def check():
+    path = "audit_logs.txt"
+    if not os.path.exists(path):
+        print("Logs not found.")
+        return
+        
+    try:
+        # PowerShell results might be in UTF-16
+        with open(path, "rb") as f:
+            raw = f.read()
+            content = raw.decode("utf-16", "ignore")
+    except:
+        content = open(path, "r", errors="ignore").read()
+        
+    questions = [
+        "What are the requirements to be listed on the Nasdaq?",
+        "When was Rule. 2010"
+    ]
+    
+    print("=== Consistency Check ===")
+    for q in questions:
+        found = q.lower() in content.lower()
+        print(f"Found '{q}': {found}")
+
+if __name__ == "__main__":
+    check()

client.py

@@ -0,0 +1,46 @@
+from openenv.core.env_client import EnvClient
+from models import ComplianceAction
+from server import ComplianceEnv
+import json
+
+def run_test():
+    print("Initializing test environment...")
+    
+    # In a real setup, client talks to a served environment via EnvClient
+    # For local demonstration, we instantiate the Env directly
+    env = ComplianceEnv()
+    
+    obs = env.reset()
+    print("Initial Observation:")
+    print(obs.model_dump_json(indent=2))
+    
+    # Simulate a bad action (approve a sanctioned entity)
+    # Let's force it to be a sanctioned entity for testing
+    env.db.current_entity = "EvilCorp"
+    
+    action1 = ComplianceAction(
+        decision="APPROVE",
+        reasoning="The transaction seems fine.",
+        cited_regulation_id="NONE"
+    )
+    
+    new_obs, reward, done, info = env.step(action1)
+    print("\nAction 1: APPROVE EvilCorp")
+    print(f"Reward: {reward} (Expect negative)")
+    print(f"Trace info: {info}")
+    
+    # Simulate a good action
+    env.db.current_entity = "EvilCorp"
+    action2 = ComplianceAction(
+        decision="BLOCK",
+        reasoning="Entity is on the sanctions list. Blocking as per AML regulations.",
+        cited_regulation_id="EU-AML-01"
+    )
+    
+    new_obs, reward, done, info = env.step(action2)
+    print("\nAction 2: BLOCK EvilCorp with good reasoning")
+    print(f"Reward: {reward} (Expect positive)")
+    print(f"Trace info: {info}")
+
+if __name__ == "__main__":
+    run_test()

client_env.py

@@ -0,0 +1,67 @@
+"""
+AegisGym WebSocket client — concrete subclass of openenv EnvClient.
+
+Usage:
+    client = AegisGymWsClient()
+    sync_client = client.sync()
+    obs = sync_client.reset()
+    obs = sync_client.step({...})
+"""
+from typing import Any, Dict
+from openenv.core.env_client import EnvClient
+from openenv.core.sync_client import SyncEnvClient
+from models import AuditAction, AuditObservation
+
+HF_SPACE_WSS = "wss://armaan020-aegisgym.hf.space"
+
+
+class AegisGymWsClient(EnvClient):
+    """Concrete EnvClient implementation for the AegisGym HF Space."""
+
+    def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]:
+        """Convert an action dict into the WS step payload."""
+        return action if isinstance(action, dict) else action.model_dump()
+
+    def _parse_result(self, payload: Dict[str, Any]) -> Any:
+        """Parse reset/step response from the server into usable result."""
+        return payload  # keep as dict; training code accesses .observation, .reward, .done
+
+    def _parse_state(self, payload: Dict[str, Any]) -> Any:
+        """Parse the state endpoint response."""
+        return payload
+
+
+def get_sync_client(ws_url: str = HF_SPACE_WSS) -> SyncEnvClient:
+    """Return a synchronous wrapper over the WebSocket client."""
+    return AegisGymWsClient(base_url=ws_url).sync()
+
+
+def test_live_env():
+    print("=== Testing Live AegisGym Space (wss) ===")
+    client = get_sync_client()
+
+    print("reset() ...")
+    result = client.reset()
+    print(f"  result type: {type(result)}")
+    print(f"  keys: {list(result.keys()) if isinstance(result, dict) else dir(result)}")
+
+    action = AuditAction(
+        action_type="FLAG",
+        target_id="ACC-BL-001",
+        regulation_citation="EU-AI-Act-Art-57"
+    ).model_dump()
+
+    print("step(FLAG ACC-BL-001) ...")
+    result = client.step(action)
+    print(f"  reward={result.get('reward') if isinstance(result, dict) else getattr(result,'reward',None)}")
+    print(f"  done={result.get('done') if isinstance(result, dict) else getattr(result,'done',None)}")
+
+    print("state() ...")
+    s = client.state()
+    print(f"  State: {s}")
+
+    print("\n=== Live environment OK! ===")
+
+
+if __name__ == "__main__":
+    test_live_env()

deploy_hf.py

@@ -0,0 +1,74 @@
+import argparse
+import subprocess
+import sys
+import os
+
+def check_huggingface_login():
+    """Verify the user is logged into Hugging Face CLI."""
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.whoami()
+        return True
+    except Exception:
+        return False
+
+def init_space(space_name: str, private: bool = True):
+    """Initialize a Hugging Face space for the environment."""
+    print(f"--- Initializing Hugging Face Space: {space_name} ---")
+    
+    if not check_huggingface_login():
+        print("Error: You are not logged in to Hugging Face CLI. Please run 'huggingface-cli login' first.")
+        sys.exit(1)
+        
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.create_repo(repo_id=space_name, repo_type="space", space_sdk="docker", private=private, exist_ok=True)
+        print("Space created or verified successfully!")
+    except Exception as e:
+        print(f"Failed to create space: {e}")
+        print("Note: If it already exists, deployment will proceed.")
+
+def deploy_to_space(space_name: str):
+    """Deploy the current directory to the Hugging Face Space."""
+    print(f"\n--- Deploying AegisGym to {space_name} ---")
+    
+    # Normally this would involve git add/commit/push to the HF remote
+    # or using the huggingface_hub Python library to upload the folder seamlessly.
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        print("Uploading files to Hub...")
+        
+        # Don't upload the cache or local virtual environments
+        ignore_patterns = ["__pycache__/*", "*.git*", ".env", "venv/*"]
+        
+        url = api.upload_folder(
+            folder_path=".",
+            repo_id=space_name,
+            repo_type="space",
+            ignore_patterns=ignore_patterns
+        )
+        print(f"Deployment successful! Your environment is live at: {url}")
+        
+    except ImportError:
+        print("Error: 'huggingface_hub' is not installed. Run 'pip install huggingface_hub'.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Deployment failed: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Deploy AegisGym to Hugging Face Spaces.")
+    parser.add_argument("space_name", help="The name of your huggingface space (e.g. username/aegisgym)")
+    parser.add_argument("--public", action="store_true", help="Make the space public (default is private)")
+    
+    args = parser.parse_args()
+    
+    init_space(args.space_name, private=not args.public)
+    deploy_to_space(args.space_name)
+    
+    print("\n--- Next Steps ---")
+    print("1. Your Space is building the Docker container using your OpenEnv configurations.")
+    print(f"2. You can access it directly via: https://huggingface.co/spaces/{args.space_name}")
+    print("3. In your training script, switch your Env URL to match your Space.")

generate_matching_report.py

@@ -0,0 +1,23 @@
+import json
+
+def generate_report():
+    with open("synced_report.json", "r") as f:
+        data = json.load(f)
+        
+    with open("matching_dataset_logs.md", "w", encoding="utf-8") as f:
+        f.write("# Synchronized Audit & Dataset Report\n\n")
+        f.write("This report pairs each simulation episode with the exact dataset entry that triggered it.\n\n")
+        
+        for e in data:
+            f.write(f"## Episode {e['episode']}\n")
+            f.write(f"### 📊 Dataset Entry\n")
+            f.write(f"**Question:** {e['dataset_question']}\n\n")
+            f.write(f"**Context (Answer):** {e['dataset_answer']}\n\n")
+            f.write(f"### 📄 Audit Log\n")
+            f.write(f"**LLM Reasoning:**\n```json\n{e['llm_reasoning']}\n```\n")
+            f.write(f"**Action:** {e['action']['action_type']} on {e['action']['target_id']}\n\n")
+            f.write(f"**Reward:** {e['reward']}\n\n")
+            f.write("---\n\n")
+
+if __name__ == "__main__":
+    generate_report()

grader.py

@@ -0,0 +1,49 @@
+import random
+from typing import Dict, Any, List
+
+class Grader:
+    def __init__(self):
+        # Dummy blacklist for Easy task
+        self.sanctioned_ids = ["ACC-BL-001", "ACC-BL-002", "ACC-BL-999"]
+    
+    def grade_easy(self, target_id: str, is_flagged: bool) -> float:
+        """
+        Sanction Hit: Binary match of AccountID vs Blacklist.
+        Returns identification accuracy (Pid)
+        """
+        is_sanctioned = target_id in self.sanctioned_ids
+        if is_sanctioned and is_flagged:
+            return 1.0
+        elif not is_sanctioned and not is_flagged:
+            return 1.0
+        return 0.0
+
+    def grade_medium(self, transactions: List[Dict[str, Any]], target_id: str, is_flagged: bool) -> float:
+        """
+        Smurfing Detection: Temporal window analysis of related transactions.
+        Checks for multiple transactions just under $10,000.
+        """
+        smurf_count = sum(1 for tx in transactions if 9000 <= tx.get("amount", 0) < 10000)
+        has_smurfing = smurf_count >= 3
+        
+        if has_smurfing and is_flagged:
+            return 1.0
+        elif not has_smurfing and not is_flagged:
+            return 1.0
+        return 0.0
+
+    def grade_hard(self, citation: str) -> float:
+        """
+        Regulatory Mapping: LLM-as-a-judge + Cosine similarity (mocked here).
+        Returns citation accuracy (Pcit) based on simulated semantic match.
+        """
+        if not citation:
+            return 0.0
+            
+        # Simplified mock evaluation
+        valid_citations = ["EU-AI-Act-Art-57", "BSA-31-USC-5318"]
+        if any(v in citation for v in valid_citations):
+            return 1.0
+        elif len(citation) > 10:
+            return 0.5 # Partial credit for trying to cite something substantive
+        return 0.0

inference.py

@@ -0,0 +1,62 @@
+"""
+AegisGym Baseline Inference Script
+Requirement: Must use OpenAI client and env vars for credentials.
+"""
+import os
+import json
+from openai import OpenAI
+from client_env import get_sync_client
+from models import AuditAction
+
+# ─── Config (Required by Meta OpenEnv) ──────────────────────────────────────
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME   = os.getenv("MODEL_NAME", "gpt-4o")
+HF_TOKEN     = os.getenv("HF_TOKEN", "") # Used if hitting HF Inference Endpoints
+ENV_URL      = os.getenv("ENV_URL", "https://armaan020-aegisgym.hf.space")
+
+client = OpenAI(api_key=HF_TOKEN or os.getenv("OPENAI_API_KEY"), base_url=API_BASE_URL)
+
+SYSTEM_PROMPT = """You are a financial compliance auditor AI.
+Respond ONLY with a JSON object:
+{"action_type": "FLAG|APPROVE|REQUEST_INFO", "target_id": "<id>", "regulation_citation": "<cite>"}"""
+
+def run_baseline(num_steps=3):
+    print(f"=== AegisGym Baseline Inference ===")
+    print(f"Model: {MODEL_NAME} | Env: {ENV_URL}")
+    
+    env = get_sync_client(ENV_URL)
+    
+    for i in range(num_steps):
+        print(f"\n--- Step {i+1} ---")
+        obs_payload = env.reset()
+        obs = obs_payload.get("observation", {})
+        
+        user_msg = (
+            f"Transactions: {obs.get('transactions', [])}\n"
+            f"Context: {obs.get('retrieved_regs', '')}\n"
+            f"Account: {obs.get('account_metadata', {})}"
+        )
+        
+        # OpenAI API Call
+        response = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_msg}
+            ],
+            response_format={"type": "json_object"}
+        )
+        
+        content = response.choices[0].message.content
+        print(f"LLM Response: {content}")
+        
+        # Parse and Step
+        try:
+            action_data = json.loads(content)
+            result = env.step(action_data)
+            print(f"Reward: {result.get('reward')} | Done: {result.get('done')}")
+        except Exception as e:
+            print(f"Parse Error: {e}")
+
+if __name__ == "__main__":
+    run_baseline()

matching_dataset_logs.md

@@ -0,0 +1,61 @@
+# Synchronized Audit & Dataset Report
+
+This report pairs each simulation episode with the exact dataset entry that triggered it.
+
+## Episode 1
+### 📊 Dataset Entry
+**Question:** When was Rule. 2010  Standards of Commercial Honor and Principles of Trade  Adopted?
+
+**Context (Answer):** 2009-05-29 00:00:00
+
+### 📄 Audit Log
+**LLM Reasoning:**
+```json
+{"action_type": "FLAG", "target_id": "ACC-BL-001", "regulation_citation": "EU-AI-Act-Art-57"}
+```
+**Action:** FLAG on ACC-BL-001
+
+**Reward:** 0.5
+
+---
+
+## Episode 2
+### 📊 Dataset Entry
+**Question:** When was Rule Rule 2020. Use of Manipulative, Deceptive or Other Fraudulent Devices adopted
+
+**Context (Answer):** 2009-05-29 00:00:00
+
+### 📄 Audit Log
+**LLM Reasoning:**
+```json
+{
+  "action_type": "FLAG",
+  "target_id": "ACC-CLEAN-01",
+  "regulation_citation": "EU-AI-Act-Art-57"
+}
+```
+**Action:** FLAG on ACC-CLEAN-01
+
+**Reward:** -1.0
+
+---
+
+## Episode 3
+### 📊 Dataset Entry
+**Question:** Which countries currently have bilateral FTAs in effect with the U.S.?
+
+
+
+**Context (Answer):** "The U.S. currently has bilateral FTAs in effect with Australia, Bahrain, Chile, Israel, Jordan, Morocco, Oman, Singapore, Peru, Korea, and Columbia.
+
+### 📄 Audit Log
+**LLM Reasoning:**
+```json
+{"action_type": "FLAG", "target_id": "ACC-BL-001", "regulation_citation": "EU-AI-Act-Art-57: Sandboxes required for high risk."}
+```
+**Action:** FLAG on ACC-BL-001
+
+**Reward:** 0.5
+
+---
+

models.py

@@ -0,0 +1,23 @@
+from openenv.core.env_server import Action, Observation, State
+from pydantic import Field
+from typing import List, Dict, Any
+
+class AuditAction(Action):
+    action_type: str = Field(description="One of: ['APPROVE', 'FLAG', 'REQUEST_INFO']")
+    target_id: str = Field(description="The ID of the account or transaction being evaluated.")
+    regulation_citation: str = Field(description="The specific regulation clause cited.")
+
+class AuditObservation(Observation):
+    transactions: List[Dict[str, Any]] = Field(description="List of transaction dicts")
+    account_metadata: Dict[str, Any] = Field(description="Account metadata")
+    retrieved_regs: str = Field(description="RAG-retrieved sections of guidelines")
+    
+    # Official OpenEnv returns these properties directly on the Observation
+    reward: float = Field(default=0.0, description="Reward gained in the step")
+    done: bool = Field(default=False, description="Whether the episode is complete")
+
+class AuditState(State):
+    step_count: int = Field(description="Current step of the episode")
+    current_tier: str = Field(description="The compliance tier active for the episode")
+    
+    # Required by base State (depending on library version, might require extra fields, but standard is dict-like or these fields)

openenv.yaml

@@ -0,0 +1,19 @@
+id: aegis_gym
+version: "1.0.0"
+name: "AegisGym: Financial Compliance & AML Sandbox"
+description: "A sandbox for training agents to perform banking audits, sanction checks, and anti-money laundering (AML) detection."
+author: "Armaan"
+tags: ["finance", "compliance", "aml", "audit", "openenv"]
+tasks:
+  - id: easy_audit
+    name: "Sanction Check"
+    description: "Evaluate a single account against a known sanctions list. Identify 'ACC-BL-*' accounts as high-risk."
+    difficulty: "easy"
+  - id: medium_audit
+    name: "Smurfing Detection"
+    description: "Analyze a stream of transactions to detect 'smurfing' (structuring) patterns where multiple deposits just under $10,000 are made."
+    difficulty: "medium"
+  - id: hard_audit
+    name: "Regulatory Alignment"
+    description: "Audit high-value transactions and cite the specific regulatory framework (e.g., EU AI Act, BSA) correctly."
+    difficulty: "hard"

requirements.txt

@@ -0,0 +1,4 @@
+openenv-core @ git+https://github.com/meta-pytorch/OpenEnv.git
+pydantic
+uvicorn
+openai

sample_dataset.py

@@ -0,0 +1,19 @@
+from datasets import load_dataset
+
+def sample_dataset():
+    print("Loading SecureFinAI-Lab/Regulations_QA...")
+    try:
+        # We'll use a streaming load or just take the first few examples
+        ds = load_dataset("SecureFinAI-Lab/Regulations_QA", split="train", streaming=True)
+        iterator = iter(ds)
+        print("\n--- Example 1 ---")
+        item1 = next(iterator)
+        print(item1)
+        print("\n--- Example 2 ---")
+        item2 = next(iterator)
+        print(item2)
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+
+if __name__ == "__main__":
+    sample_dataset()

server.py

@@ -0,0 +1,88 @@
+import random
+from typing import Dict, Any, Tuple
+from openenv.core.env_server import Environment
+from models import AuditAction, AuditObservation, AuditState
+from grader import Grader
+
+class AegisGymEnv(Environment):
+    def __init__(self):
+        super().__init__()
+        self.grader = Grader()
+        self.step_count = 0
+        self.current_tier = "easy"
+        self.max_steps = 100
+        
+        # Test Data
+        self.current_target_id = "ACC-BL-001"
+        self.current_transactions = []
+        
+    def reset(self) -> AuditObservation:
+        self.step_count = 0
+        self._set_next_scenario()
+        return self._create_observation(0.0, False)
+
+    @property
+    def state(self) -> AuditState:
+        return AuditState(
+            step_count=self.step_count,
+            current_tier=self.current_tier
+        )
+        
+    def _create_observation(self, reward: float, done: bool) -> AuditObservation:
+        return AuditObservation(
+            transactions=self.current_transactions,
+            account_metadata={"age_days": 120, "tier": "standard", "target_id": self.current_target_id},
+            retrieved_regs="EU-AI-Act-Art-57: Sandboxes required for high risk. BSA-31-USC-5318: AML programs.",
+            reward=reward,
+            done=done
+        )
+
+    def step(self, action: AuditAction) -> AuditObservation:
+        self.step_count += 1
+        
+        is_flagged = action.action_type in ["FLAG", "BLOCK", "REQUEST_INFO"]
+        
+        p_id, p_cit = 0.0, 0.0
+        f_p, f_n = 0.0, 0.0
+        
+        if self.current_tier == "easy":
+            p_id = self.grader.grade_easy(action.target_id, is_flagged)
+            is_sanctioned = action.target_id in self.grader.sanctioned_ids
+            if not is_sanctioned and is_flagged: f_p = 1.0
+            elif is_sanctioned and not is_flagged: f_n = 5.0
+                
+        elif self.current_tier == "medium":
+            p_id = self.grader.grade_medium(self.current_transactions, action.target_id, is_flagged)
+            smurf_count = sum(1 for tx in self.current_transactions if 9000 <= tx.get("amount", 0) < 10000)
+            has_smurfing = smurf_count >= 3
+            if not has_smurfing and is_flagged: f_p = 1.0
+            elif has_smurfing and not is_flagged: f_n = 5.0
+                
+        elif self.current_tier == "hard":
+            p_id = 1.0 if is_flagged else 0.0
+            p_cit = self.grader.grade_hard(action.regulation_citation)
+            if not is_flagged: f_n = 5.0
+
+        reward = (0.5 * p_id) + (0.5 * p_cit) - (1.0 * f_p) - (1.0 * f_n)
+        done = self.step_count >= self.max_steps
+        
+        obs = self._create_observation(reward, done)
+        self._set_next_scenario()
+        return obs
+
+    def _set_next_scenario(self):
+        tiers = ["easy", "medium", "hard"]
+        self.current_tier = tiers[self.step_count % 3]
+        
+        if self.current_tier == "easy":
+            self.current_target_id = random.choice(["ACC-BL-001", "ACC-CLEAN-01"])
+            self.current_transactions = [{"amount": 500, "currency": "USD"}]
+        elif self.current_tier == "medium":
+            self.current_target_id = "ACC-SMURF-99"
+            if random.random() > 0.5:
+                self.current_transactions = [{"amount": 9500, "currency": "USD"} for _ in range(4)]
+            else:
+                self.current_transactions = [{"amount": 5000, "currency": "USD"}, {"amount": 200, "currency": "USD"}]
+        elif self.current_tier == "hard":
+            self.current_target_id = "ACC-REG-VIOLATOR"
+            self.current_transactions = [{"amount": 50000, "currency": "USD", "note": "High risk AI deployment without sandbox"}]

simulation.py

@@ -0,0 +1,82 @@
+"""
+AegisGym Simulation Script
+Runs multiple audit episodes using the LLM for inference (CPU-friendly).
+This generates the logs and metrics to analyze the system's performance.
+"""
+import os
+import torch
+import json
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from client_env import get_sync_client
+from train import parse_action, SYSTEM_PROMPT
+
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+ENV_URL = "https://armaan020-aegisgym.hf.space"
+
+def run_simulation(num_episodes=5):
+    print(f"=== Starting AegisGym Simulation (Inference Only) ===")
+    print(f"Model: {MODEL_NAME}")
+    print(f"Env:   {ENV_URL}\n")
+
+    print(f"Loading model on CPU...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype="auto", device_map="cpu")
+    print("Model loaded.\n")
+
+    env = get_sync_client(ENV_URL)
+    
+    total_reward = 0
+    results = []
+
+    for i in range(num_episodes):
+        print(f"--- Episode {i+1} ---")
+        result = env.reset()
+        obs_dict = result.get("observation", {})
+        
+        state = env.state()
+        tier = state.get("current_tier", "easy")
+        
+        user_msg = (
+            f"Audit the following transaction.\n\n"
+            f"Tier: {tier.upper()}\n"
+            f"Transactions: {obs_dict.get('transactions', [])}\n"
+            f"Context: {obs_dict.get('retrieved_regs', [])}\n"
+            f"Account: {obs_dict.get('account_metadata', {})}"
+        )
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user",   "content": user_msg},
+        ]
+        prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        
+        print(f"[Audit Prompt]:\n{user_msg}")
+        
+        # Inference
+        inputs = tokenizer(prompt, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=128)
+        
+        completion = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        print(f"[Model Reasoning]:\n{completion}")
+        
+        action = parse_action(completion)
+        print(f"[Action]: {action.action_type} on {action.target_id}")
+        
+        step_result = env.step(action.model_dump())
+        reward = step_result.get("reward", 0.0)
+        done = step_result.get("done", False)
+        
+        print(f"[Reward]: {reward} | [Done]: {done}\n")
+        total_reward += reward
+        results.append({
+            "episode": i+1,
+            "tier": tier,
+            "action": action.action_type,
+            "reward": reward
+        })
+
+    print(f"=== Simulation Complete ===")
+    print(f"Average Reward: {total_reward / num_episodes}")
+
+if __name__ == "__main__":
+    run_simulation(num_episodes=3)

synced_report.json

@@ -0,0 +1,44 @@
+[
+  {
+    "episode": 1,
+    "dataset_question": "When was Rule. 2010  Standards of Commercial Honor and Principles of Trade  Adopted?",
+    "dataset_answer": "2009-05-29 00:00:00",
+    "tier": "easy",
+    "llm_reasoning": "{\"action_type\": \"FLAG\", \"target_id\": \"ACC-BL-001\", \"regulation_citation\": \"EU-AI-Act-Art-57\"}",
+    "action": {
+      "metadata": {},
+      "action_type": "FLAG",
+      "target_id": "ACC-BL-001",
+      "regulation_citation": "EU-AI-Act-Art-57"
+    },
+    "reward": 0.5
+  },
+  {
+    "episode": 2,
+    "dataset_question": "When was Rule Rule 2020. Use of Manipulative, Deceptive or Other Fraudulent Devices adopted",
+    "dataset_answer": "2009-05-29 00:00:00",
+    "tier": "easy",
+    "llm_reasoning": "{\n  \"action_type\": \"FLAG\",\n  \"target_id\": \"ACC-CLEAN-01\",\n  \"regulation_citation\": \"EU-AI-Act-Art-57\"\n}",
+    "action": {
+      "metadata": {},
+      "action_type": "FLAG",
+      "target_id": "ACC-CLEAN-01",
+      "regulation_citation": "EU-AI-Act-Art-57"
+    },
+    "reward": -1.0
+  },
+  {
+    "episode": 3,
+    "dataset_question": "Which countries currently have bilateral FTAs in effect with the U.S.?\n\n",
+    "dataset_answer": "\"The U.S. currently has bilateral FTAs in effect with Australia, Bahrain, Chile, Israel, Jordan, Morocco, Oman, Singapore, Peru, Korea, and Columbia.",
+    "tier": "easy",
+    "llm_reasoning": "{\"action_type\": \"FLAG\", \"target_id\": \"ACC-BL-001\", \"regulation_citation\": \"EU-AI-Act-Art-57: Sandboxes required for high risk.\"}",
+    "action": {
+      "metadata": {},
+      "action_type": "FLAG",
+      "target_id": "ACC-BL-001",
+      "regulation_citation": "EU-AI-Act-Art-57: Sandboxes required for high risk."
+    },
+    "reward": 0.5
+  }
+]

synced_simulation.py

@@ -0,0 +1,75 @@
+"""
+AegisGym Synchronized Simulation
+Saves exact dataset entries alongside the audit logs they generated.
+"""
+import torch
+import json
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import load_dataset
+from client_env import get_sync_client
+from train import parse_action, SYSTEM_PROMPT
+
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+ENV_URL = "https://armaan020-aegisgym.hf.space"
+
+def run_synced_simulation(num_episodes=3):
+    print(f"Loading model on CPU...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype="auto", device_map="cpu")
+    
+    print("Loading dataset...")
+    ds = load_dataset("SecureFinAI-Lab/Regulations_QA", split="train", streaming=True)
+    it = iter(ds)
+    
+    env = get_sync_client(ENV_URL)
+    full_report = []
+
+    for i in range(num_episodes):
+        print(f"--- Episode {i+1} ---")
+        item = next(it)
+        result = env.reset()
+        obs_dict = result.get("observation", {})
+        state = env.state()
+        tier = state.get("current_tier", "easy")
+        
+        custom_prompt = item.get("question", "Audit the following transaction.")
+        dataset_answer = item.get("answer", "No specific guidance provided.")
+        
+        user_msg = (
+            f"{custom_prompt}\n\n"
+            f"Tier: {tier.upper()}\n"
+            f"Transactions: {obs_dict.get('transactions', [])}\n"
+            f"Context: {obs_dict.get('retrieved_regs', [])}\n"
+            f"Regulatory Hint: {dataset_answer}\n"
+            f"Account: {obs_dict.get('account_metadata', {})}"
+        )
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user",   "content": user_msg},
+        ]
+        prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        
+        inputs = tokenizer(prompt, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=128)
+        
+        completion = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        action = parse_action(completion)
+        step_result = env.step(action.model_dump())
+        
+        full_report.append({
+            "episode": i+1,
+            "dataset_question": custom_prompt,
+            "dataset_answer": dataset_answer,
+            "tier": tier,
+            "llm_reasoning": completion,
+            "action": action.model_dump(),
+            "reward": step_result.get("reward", 0.0)
+        })
+
+    with open("synced_report.json", "w") as f:
+        json.dump(full_report, f, indent=2)
+    print("\nSaved synced_report.json")
+
+if __name__ == "__main__":
+    run_synced_simulation()

train.py

@@ -0,0 +1,155 @@
+"""
+AegisGym GRPO Training Script
+Connects to the live HF Space for environment rollouts.
+Run: python train.py
+"""
+import os
+import json
+from datasets import Dataset
+from trl import GRPOConfig, GRPOTrainer
+from transformers import AutoTokenizer
+from client_env import get_sync_client
+from models import AuditAction
+
+# ─── Config ──────────────────────────────────────────────────────────────────
+ENV_URL    = os.getenv("AEGISGYM_URL", "https://armaan020-aegisgym.hf.space")
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-0.5B-Instruct")
+OUTPUT_DIR = "aegisgym-grpo-agent"
+
+SYSTEM_PROMPT = """You are a financial compliance auditor AI.
+Given a transaction scenario and regulatory context, respond with a JSON object:
+{"action_type": "FLAG|APPROVE|REQUEST_INFO", "target_id": "<account_id>", "regulation_citation": "<regulation>"}
+Be precise and concise."""
+
+# ─── Action parser ────────────────────────────────────────────────────────────
+def parse_action(text: str) -> AuditAction:
+    try:
+        start = text.find("{")
+        end   = text.rfind("}") + 1
+        if start >= 0 and end > start:
+            data = json.loads(text[start:end])
+            return AuditAction(**data)
+    except Exception:
+        pass
+    return AuditAction(action_type="REQUEST_INFO", target_id="UNKNOWN", regulation_citation="parse_error")
+
+# ─── Dataset Loading ──────────────────────────────────────────────────────────
+from datasets import load_dataset
+import itertools
+
+print("Loading SecureFinAI-Lab/Regulations_QA dataset...")
+try:
+    # Use a streaming dataset for efficiency
+    raw_dataset = load_dataset("SecureFinAI-Lab/Regulations_QA", split="train", streaming=True)
+    dataset_iterator = itertools.cycle(iter(raw_dataset))
+    print("Dataset loaded successfully.")
+except Exception as e:
+    print(f"Warning: Failed to load dataset: {e}. Falling back to default prompts.")
+    dataset_iterator = None
+
+# ─── Rollout function ─────────────────────────────────────────────────────────
+def rollout_func(trainer, prompts, tokenizer):
+    """One rollout episode connecting to the live AegisGym Space with dataset augmentation."""
+    from trl.experimental.openenv import generate_rollout_completions
+    
+    env = get_sync_client(ENV_URL)
+    result = env.reset()
+    obs_dict = result.get("observation", {})
+    
+    # Sample from dataset if available
+    dataset_context = ""
+    custom_prompt = "Audit the following transaction."
+    if dataset_iterator:
+        item = next(dataset_iterator)
+        custom_prompt = item.get("question", custom_prompt)
+        dataset_context = f"\nRegulatory Context: {item.get('answer', '')}"
+    
+    all_prompt_ids, all_completion_ids, all_logprobs, rewards = [], [], [], []
+    
+    state = env.state()
+    tier = state.get("current_tier", "easy")
+    
+    # Combine dataset context with environment context
+    regs = obs_dict.get('retrieved_regs', "")
+    retrieved_regs = [regs] if isinstance(regs, str) else list(regs)
+    if dataset_context:
+        retrieved_regs.append(dataset_context)
+        
+    user_msg = (
+        f"{custom_prompt}\n\n"
+        f"Tier: {tier.upper()}\n"
+        f"Transactions: {obs_dict.get('transactions', [])}\n"
+        f"Context: {retrieved_regs}\n"
+        f"Account: {obs_dict.get('account_metadata', {})}"
+    )
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user",   "content": user_msg},
+    ]
+    prompt_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    
+    out = generate_rollout_completions(trainer, [prompt_text])[0]
+    all_prompt_ids.extend(out["prompt_ids"])
+    all_completion_ids.extend(out["completion_ids"])
+    all_logprobs.extend(out["logprobs"])
+    
+    completion_text = out.get("text") or tokenizer.decode(out["completion_ids"], skip_special_tokens=True)
+    action = parse_action(completion_text)
+    result = env.step(action.model_dump())
+    rewards.append(float(result.get("reward", 0.0)))
+    
+    return {
+        "prompt_ids":      all_prompt_ids,
+        "completion_ids":  all_completion_ids,
+        "logprobs":        all_logprobs,
+        "env_reward":      rewards[-1] if rewards else 0.0,
+    }
+
+# ─── Reward shim ─────────────────────────────────────────────────────────────
+def reward_compliance(completions, **kwargs):
+    rewards = kwargs.get("env_reward", [])
+    if not rewards:
+        return [0.0] * len(completions)
+    return [float(r) for r in rewards]
+
+# ─── Main ─────────────────────────────────────────────────────────────────────
+def main():
+    print(f"Model : {MODEL_NAME}")
+    print(f"Env   : {ENV_URL}")
+    
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    dataset = Dataset.from_dict({"prompt": ["Audit the following transaction."] * 200})
+    
+    config = GRPOConfig(
+        num_train_epochs              = 1,
+        learning_rate                 = 5e-6,
+        per_device_train_batch_size   = 1,
+        gradient_accumulation_steps   = 4,
+        warmup_steps                  = 10,
+        num_generations               = 2,
+        max_completion_length         = 256,
+        use_vllm                      = False,
+        output_dir                    = OUTPUT_DIR,
+        logging_steps                 = 1,
+        save_steps                    = 25,
+        gradient_checkpointing        = True,
+    )
+    
+    trainer = GRPOTrainer(
+        model            = MODEL_NAME,
+        processing_class = tokenizer,
+        reward_funcs     = [reward_compliance],
+        train_dataset    = dataset,
+        args             = config,
+        rollout_func     = rollout_func,
+    )
+    
+    print("\n=== Starting GRPO Training ===")
+    trainer.train()
+    trainer.save_model(OUTPUT_DIR)
+    print(f"\nModel saved to {OUTPUT_DIR}/")
+
+if __name__ == "__main__":
+    main()

verify_consistency.py

@@ -0,0 +1,26 @@
+from datasets import load_dataset
+import os
+
+def verify():
+    log_path = "audit_logs.txt"
+    if not os.path.exists(log_path):
+        print("Logs not found.")
+        return
+        
+    # Read logs (UTF-16 from PowerShell redirection)
+    with open(log_path, "rb") as f:
+        content = f.read().decode("utf-16", "ignore")
+        
+    print("--- Loading Dataset ---")
+    ds = load_dataset("SecureFinAI-Lab/Regulations_QA", split="train", streaming=True)
+    it = iter(ds)
+    
+    print("--- Comparing Samples ---")
+    for i in range(5):
+        sample = next(it)
+        q = sample["question"]
+        found = q in content
+        print(f"Sample {i+1}: '{q[:50]}...' Found: {found}")
+
+if __name__ == "__main__":
+    verify()

verify_dataset_rollout.py

@@ -0,0 +1,27 @@
+"""
+Verification script for dataset-augmented rollout function.
+"""
+from transformers import AutoTokenizer
+from train import rollout_func, dataset_iterator
+import trl.experimental.openenv
+
+def verify():
+    print("=== Verifying Dataset-Augmented Rollout ===")
+    model_name = "Qwen/Qwen2.5-0.5B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+    # Mock generation
+    def mock_gen(trainer, prompts):
+        print(f"\n[Generated Prompt Sample]:\n{prompts[0][:500]}...\n")
+        fake_completion = '{"action_type": "APPROVE", "target_id": "ACC-001", "regulation_citation": "EU-Regulation"}'
+        return [{"prompt_ids": [1], "completion_ids": [1], "logprobs": [0.0], "text": fake_completion}]
+    
+    trl.experimental.openenv.generate_rollout_completions = mock_gen
+    
+    print("Running rollout_func...")
+    result = rollout_func(None, ["prompt"], tokenizer)
+    print(f"\nRollout Finished!")
+    print(f"Reward: {result['env_reward']}")
+
+if __name__ == "__main__":
+    verify()