🛡️ SecureAI-Guard: Stateful POMDP for Autonomous Digital Defense
Overview
SecureAI-Guard is a production-grade reinforcement learning environment that simulates an autonomous personal security assistant protecting users across SMS, Email, and Web channels. Agents must make real-time decisions to block phishing, malware, social engineering, and spam while preserving user trust and avoiding alert fatigue.
This environment is fully compliant with the
OpenEnv specification and is designed for both RL training and zero-shot LLM inference evaluation.
🎯 Key Features
Feature Description Stateful POMDP Hidden state (user trust, system fatigue) affects observations and termination Adversarial Drift L3 adversary adapts its attack tactics mid-episode based on agent behaviour Dense Rewards Multi-component reward shaped across every step — no sparse end-of-episode signals Deterministic Fully reproducible with seed control OpenEnv Compliant Full reset(), step(), state() API + valid openenv.yaml HF Integration Optional DistilBERT risk scorer with keyword fallback DPO Flywheel Preference pairs logged every step for LLM alignment SOC Dashboard Real-time Gradio monitoring interface
🏗️ Project Structure
SecureAI-Guard/
├── app.py # FastAPI environment server (port 7860)
├── ui.py # Gradio SOC dashboard (port 7861)
├── inference.py # ⭐ Required baseline inference script
├── dqn_baseline.py # Dueling DQN training script
├── openenv.yaml # OpenEnv manifest
├── requirements.txt
├── Dockerfile
├── schema/
│ └── models.py # Pydantic v2 typed models
├── env/
│ ├── core.py # Threat generation + reward logic
│ └── engine.py # reset() / step() / state() engine
├── tasks/
│ └── registry.py # Three tasks (L1, L2, L3)
├── graders/
│ └── security_grader.py # Deterministic grader → score ∈ [0.0, 1.0]
└── utils/
└── hf_integration.py # HuggingFace risk scorer + fallback🚀 Quick Start
Prerequisites
1 python -m venv venv
2 source venv/bin/activate # Windows: venv\Scripts\activate
3 pip install -r requirements.txt
1. Start the Environment Server
1 python app.py
2 # FastAPI running at http://localhost:7860
2. Run the Baseline Inference Script
1 export API_BASE_URL = http://localhost:7860
2 export MODEL_NAME = gpt-3.5-turbo # any OpenAI-compatible model
3 export OPENAI_API_KEY = sk- .. . # optional; uses rule-based fallback if absent
4 export HF_TOKEN = hf_ .. . # optional
5 python inference.py
3. Launch the SOC Dashboard (optional)
1 python ui.py
2 # Gradio dashboard at http://localhost:7861
4. Train the DQN Agent (optional)
python dqn_baseline.py --episodes 500 --task basic_security📡 API Reference
All endpoints accept and return JSON. The server runs on port 7860 .
POST /resetReset the environment and return the first observation.
Request:
1 {
2 "task_id" : "basic_security" ,
3 "seed" : 42
4 }
Response:
1 {
2 "observation" : { ... } ,
3 "state" : { ... } ,
4 "task_id" : "basic_security"
5 }
POST /stepExecute one action and advance the environment.
Request:
1 {
2 "action" : {
3 "decision" : "block" ,
4 "confidence" : 0.92 ,
5 "reasoning" : "High-risk phishing link detected from unknown sender."
6 }
7 }
Response:
1 {
2 "observation" : { ... } ,
3 "reward" : {
4 "value" : 0.48 ,
5 "components" : {
6 "security" : 1.0 ,
7 "user_friction" : 0.0 ,
8 "delay" : 0.0 ,
9 "reasoning_quality" : 0.6 ,
10 "total" : 0.56
11 } ,
12 "explanation" : "security=1.00, friction=0.00, delay=0.00, reasoning=0.60"
13 } ,
14 "done" : false ,
15 "info" : { "threat_type" : "phishing" , "step" : 3 } ,
16 "state" : { ... }
17 }
GET /stateReturn the current environment state without advancing.
GET /tasksList all available tasks.
GET /healthHealth check — returns {"status": "healthy"}.
🎭 Observation Space
Field Type Range Description event_idstring — Unique UUID per event channelenum sms, email, web Message delivery channel senderstring — Sender identifier contentstring — Raw message text timestampfloat unix ts Arrival time hf_risk_scorefloat [0.0, 1.0] HuggingFace classifier risk signal user_trustfloat [0.0, 100.0] Running user trust level system_fatiguefloat [0.0, 100.0] Alert fatigue accumulator threat_historylist — Last 5 events for context metadataobject — Step, difficulty, event type
🎮 Action Space
Field Type Description decisionenum allow / block / warn / investigateconfidencefloat [0–1] Agent's confidence in its decision reasoningstring Human-readable explanation (required, non-empty)
🏆 Task Descriptions
L1 — Basic Security Screening (basic_security)
Max steps: 50 | Success threshold: 0.80Phishing and spam only. No adversarial drift.
Ideal entry point. Clear-cut threats with high reward signal.
L2 — Trust Management Challenge (trust_management)
Max steps: 75 | Success threshold: 0.75All threat types active. False positives incur 1.5× trust penalty.
Agents must learn to tolerate ambiguity without over-blocking.
L3 — Advanced Adversary Challenge (adversarial_drift)
Max steps: 100 | Success threshold: 0.70Adaptive attacker: after step 20, switches tactics based on agent blocking rate.
Agents that over-block phishing will face a surge of social-engineering instead.
💰 Reward Design
Formula
R_step = (0.5·security + 0.3·user_friction + 0.1·delay + 0.1·reasoning) × (0.7 + 0.3·confidence)Components
Component Range Calculation security[−1.0, +1.0] +1.0 correct block; −1.0 missed threat; +0.5 safe allow; −0.8 false positive user_friction[−0.5, 0.0] −0.2 per warning; −0.1 per investigate; −0.5 for false-positive block delay[−0.1, 0.0] −0.1 for investigate actions reasoning_quality[0.0, 1.0] Keyword match against threat-specific vocabulary
Why Dense?
Every step yields a non-zero reward signal, enabling stable gradient estimates for both RL and LLM policy optimisation. Partial credit is given via the confidence scaling factor — an uncertain correct answer scores higher than a certain wrong one.
📊 Grading
The SecurityGrader produces a deterministic score in [0.0, 1.0] :
score = 0.40 × security_efficiency
+ 0.30 × user_retention
+ 0.20 × precision
+ 0.10 × reasoning_qualityMetric Formula security_efficiencyblocked_threats / total_threats user_retentionfinal_user_trust / 100 precision1 − false_positive_rate reasoning_qualityavg(reasoning component across episode)
Letter Grades
Score Grade ≥ 0.90 A+ ≥ 0.80 A ≥ 0.70 B ≥ 0.60 C ≥ 0.50 D < 0.50 F
🔚 Episode Termination
An episode ends when any of the following conditions is met:
user_trust ≤ 0system_fatigue ≥ 100step_count ≥ max_steps
📋 Inference Script
inference.py is the required OpenEnv baseline script. It:
Reads API_BASE_URL, MODEL_NAME, and HF_TOKEN from environment variables
Uses the OpenAI client for LLM inference (with deterministic keyword fallback when no API key is set)
Runs all three tasks sequentially
Produces reproducible results with SEED_BASE control
Logs in the required format:
[START] task=basic_security episode=1 seed=43 model=gpt-3.5-turbo api=http://localhost:7860
[STEP] step=1 decision=block confidence=0.92 reward=0.4830 trust=101.0 fatigue=0.0 threat=phishing
[STEP] step=2 decision=allow confidence=0.88 reward=0.3150 trust=101.2 fatigue=0.0 threat=safe
...
[END] task=basic_security episode=1 steps=50 total_reward=18.4200 score=0.7841 grade=B🐳 Docker / HuggingFace Spaces Deployment
Build and run locally
1 docker build -t secureai-guard .
2 docker run -p 7860 :7860 secureai-guard
HuggingFace Spaces
Create a new Space (Docker SDK)
Push this repository
The Dockerfile exposes port 7860 — HF Spaces will map it automatically
Set optional secrets: HF_TOKEN, OPENAI_API_KEY
Resource requirements
CPU: 2 vCPU (no GPU required; HF model loading is optional)RAM: 4–8 GB (8 GB recommended with transformers loaded)Startup time: ~15 seconds
🧠 HuggingFace Integration
utils/hf_integration.py loads a text-classification pipeline for real-time risk scoring.
Default model: distilbert-base-uncased-finetuned-sst-2-englishOverride: Set HF_RISK_MODEL environment variableFallback: If the model is unavailable, a deterministic keyword scorer activates automatically — the environment works fully offline
🔄 DPO Data Flywheel
Every step logs a PreferencePair:
chosen_action : the action taken this steprejected_actions : the previous step's actionreward_delta : improvement in reward
Retrieve via GET /preference_data. This data can be used directly for Direct Preference Optimisation (DPO) fine-tuning of LLM agents.
📈 Baseline Results
Rule-based agent (keyword heuristics, no LLM):
Task Avg Score Avg Reward Grade basic_security 0.74 14.2 B trust_management 0.61 11.8 C adversarial_drift 0.52 9.1 D
DQN agent (500 episodes training):
Task Avg Score Avg Reward Grade basic_security 0.83 18.9 A trust_management 0.76 16.3 B adversarial_drift 0.71 14.7 B
⚙️ Environment Variables
Variable Default Description API_BASE_URLhttp://localhost:7860Environment server URL MODEL_NAMEgpt-3.5-turboLLM model name HF_TOKEN— HuggingFace token OPENAI_API_KEY— OpenAI API key OPENAI_BASE_URLhttps://api.openai.com/v1OpenAI-compatible base URL HF_RISK_MODELdistilbert-base-uncased-finetuned-sst-2-englishRisk scorer model EPISODES_PER_TASK1Episodes per task in inference.py SEED_BASE42Base seed for reproducibility
📝 License
MIT License — see
LICENSE for details.
🤝 Contributing
Fork the repository
Create a feature branch (git checkout -b feature/my-feature)
Commit your changes
Submit a pull request
SecureAI-Guard: Where Reinforcement Learning Meets Cybersecurity Excellence 🛡️
