secenvY 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (80) hide show
  1. secenvy-0.1.0/PKG-INFO +129 -0
  2. secenvy-0.1.0/README.md +88 -0
  3. secenvy-0.1.0/pyproject.toml +77 -0
  4. secenvy-0.1.0/secenvY/__init__.py +43 -0
  5. secenvy-0.1.0/secenvY/agents/__init__.py +17 -0
  6. secenvy-0.1.0/secenvY/agents/dqn_agent.py +225 -0
  7. secenvy-0.1.0/secenvY/agents/networks.py +81 -0
  8. secenvy-0.1.0/secenvY/agents/replay_buffer.py +205 -0
  9. secenvy-0.1.0/secenvY/core/__init__.py +109 -0
  10. secenvy-0.1.0/secenvY/core/exceptions.py +41 -0
  11. secenvy-0.1.0/secenvY/core/schemas.py +489 -0
  12. secenvy-0.1.0/secenvY/core/vocabularies.py +106 -0
  13. secenvy-0.1.0/secenvY/data/chroma/1d89f56b-3cb2-42ed-bd6f-d7f3dbf9081f/data_level0.bin +0 -0
  14. secenvy-0.1.0/secenvY/data/chroma/1d89f56b-3cb2-42ed-bd6f-d7f3dbf9081f/header.bin +0 -0
  15. secenvy-0.1.0/secenvY/data/chroma/1d89f56b-3cb2-42ed-bd6f-d7f3dbf9081f/length.bin +0 -0
  16. secenvy-0.1.0/secenvY/data/chroma/1d89f56b-3cb2-42ed-bd6f-d7f3dbf9081f/link_lists.bin +0 -0
  17. secenvy-0.1.0/secenvY/data/chroma/6f46a7d2-742d-49db-8b79-8a54d292c4cd/data_level0.bin +0 -0
  18. secenvy-0.1.0/secenvY/data/chroma/6f46a7d2-742d-49db-8b79-8a54d292c4cd/header.bin +0 -0
  19. secenvy-0.1.0/secenvY/data/chroma/6f46a7d2-742d-49db-8b79-8a54d292c4cd/length.bin +0 -0
  20. secenvy-0.1.0/secenvY/data/chroma/6f46a7d2-742d-49db-8b79-8a54d292c4cd/link_lists.bin +0 -0
  21. secenvy-0.1.0/secenvY/data/chroma/chroma.sqlite3 +0 -0
  22. secenvy-0.1.0/secenvY/data/knowledge/components.txt +61 -0
  23. secenvy-0.1.0/secenvY/data/knowledge/crypto.txt +27 -0
  24. secenvy-0.1.0/secenvY/data/knowledge/permissions.txt +64 -0
  25. secenvy-0.1.0/secenvY/data/knowledge/secrets.txt +76 -0
  26. secenvy-0.1.0/secenvY/data/knowledge/security_basics.txt +19 -0
  27. secenvy-0.1.0/secenvY/data/knowledge/storage.txt +66 -0
  28. secenvy-0.1.0/secenvY/data/knowledge/webview.txt +28 -0
  29. secenvy-0.1.0/secenvY/data/report.json +421 -0
  30. secenvy-0.1.0/secenvY/data/sample_observation.json +333 -0
  31. secenvy-0.1.0/secenvY/env/__init__.py +5 -0
  32. secenvy-0.1.0/secenvY/env/security_env.py +282 -0
  33. secenvy-0.1.0/secenvY/gateway/__init__.py +11 -0
  34. secenvy-0.1.0/secenvY/gateway/base.py +59 -0
  35. secenvy-0.1.0/secenvY/gateway/mock_gateway.py +454 -0
  36. secenvy-0.1.0/secenvY/gateway/websocket_gateway.py +288 -0
  37. secenvy-0.1.0/secenvY/services/__init__.py +26 -0
  38. secenvy-0.1.0/secenvY/services/action_mapper.py +119 -0
  39. secenvy-0.1.0/secenvY/services/episode_policy.py +74 -0
  40. secenvy-0.1.0/secenvY/services/reward_calculator.py +125 -0
  41. secenvy-0.1.0/secenvY/services/security_reasoner/__init__.py +46 -0
  42. secenvy-0.1.0/secenvY/services/security_reasoner/build_knowledge.py +60 -0
  43. secenvy-0.1.0/secenvY/services/security_reasoner/chunker.py +23 -0
  44. secenvy-0.1.0/secenvY/services/security_reasoner/config.py +53 -0
  45. secenvy-0.1.0/secenvY/services/security_reasoner/embeddings.py +18 -0
  46. secenvy-0.1.0/secenvY/services/security_reasoner/llm.py +102 -0
  47. secenvy-0.1.0/secenvY/services/security_reasoner/main.py +133 -0
  48. secenvy-0.1.0/secenvY/services/security_reasoner/rag.py +79 -0
  49. secenvy-0.1.0/secenvY/services/security_reasoner/schemas.py +9 -0
  50. secenvy-0.1.0/secenvY/services/security_reasoner/service.py +288 -0
  51. secenvy-0.1.0/secenvY/services/security_reasoner/vector_store.py +58 -0
  52. secenvy-0.1.0/secenvY/services/security_reasoner_main.py +43 -0
  53. secenvy-0.1.0/secenvY/services/state_builder.py +919 -0
  54. secenvy-0.1.0/secenvY/services/state_builder_main.py +77 -0
  55. secenvy-0.1.0/secenvY/services/state_encoder.py +379 -0
  56. secenvy-0.1.0/secenvY/services/state_encoder_main.py +84 -0
  57. secenvy-0.1.0/secenvY/services/state_validator.py +267 -0
  58. secenvy-0.1.0/secenvY/services/state_validator_main.py +59 -0
  59. secenvy-0.1.0/secenvY/services/tensor_converter.py +99 -0
  60. secenvy-0.1.0/secenvY.egg-info/PKG-INFO +129 -0
  61. secenvy-0.1.0/secenvY.egg-info/SOURCES.txt +78 -0
  62. secenvy-0.1.0/secenvY.egg-info/dependency_links.txt +1 -0
  63. secenvy-0.1.0/secenvY.egg-info/requires.txt +20 -0
  64. secenvy-0.1.0/secenvY.egg-info/top_level.txt +1 -0
  65. secenvy-0.1.0/setup.cfg +4 -0
  66. secenvy-0.1.0/tests/test_action_mapper.py +76 -0
  67. secenvy-0.1.0/tests/test_dueling_dqn.py +199 -0
  68. secenvy-0.1.0/tests/test_episode_policy.py +72 -0
  69. secenvy-0.1.0/tests/test_gymnasium_compliance.py +84 -0
  70. secenvy-0.1.0/tests/test_mock_gateway.py +149 -0
  71. secenvy-0.1.0/tests/test_pipeline_e2e.py +199 -0
  72. secenvy-0.1.0/tests/test_pipeline_performance.py +117 -0
  73. secenvy-0.1.0/tests/test_reward_calculator.py +278 -0
  74. secenvy-0.1.0/tests/test_security_reasoner.py +207 -0
  75. secenvy-0.1.0/tests/test_state_builder.py +214 -0
  76. secenvy-0.1.0/tests/test_state_encoder.py +125 -0
  77. secenvy-0.1.0/tests/test_state_validator.py +201 -0
  78. secenvy-0.1.0/tests/test_tensor_converter.py +67 -0
  79. secenvy-0.1.0/tests/test_vocabularies.py +103 -0
  80. secenvy-0.1.0/tests/test_websocket_gateway.py +153 -0
secenvy-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,129 @@
1
+ Metadata-Version: 2.4
2
+ Name: secenvY
3
+ Version: 0.1.0
4
+ Summary: Gymnasium Environment for AI-driven Android Security Analysis, Dynamic Fuzzing, and Penetration Testing
5
+ Author: Android Security Agent Team
6
+ License: Apache-2.0
7
+ Project-URL: Homepage, https://github.com/d4em0ny/android-security-agent
8
+ Project-URL: Repository, https://github.com/d4em0ny/android-security-agent.git
9
+ Project-URL: Issues, https://github.com/d4em0ny/android-security-agent/issues
10
+ Keywords: reinforcement-learning,gymnasium,android,security,penetration-testing,fuzzing,rag,dueling-dqn
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Topic :: Security
15
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3.10
18
+ Classifier: Programming Language :: Python :: 3.11
19
+ Classifier: Programming Language :: Python :: 3.12
20
+ Requires-Python: >=3.10
21
+ Description-Content-Type: text/markdown
22
+ Requires-Dist: gymnasium>=0.29.0
23
+ Requires-Dist: numpy>=1.24.0
24
+ Requires-Dist: pydantic>=2.0.0
25
+ Requires-Dist: pydantic-settings>=2.0.0
26
+ Requires-Dist: torch>=2.0.0
27
+ Requires-Dist: chromadb>=0.5.0
28
+ Requires-Dist: google-genai>=0.1.0
29
+ Requires-Dist: sentence-transformers>=2.2.0
30
+ Requires-Dist: python-dotenv>=1.0.0
31
+ Requires-Dist: rich>=13.0.0
32
+ Requires-Dist: fastapi>=0.100.0
33
+ Requires-Dist: uvicorn>=0.20.0
34
+ Requires-Dist: websockets>=11.0
35
+ Provides-Extra: dev
36
+ Requires-Dist: pytest>=7.0.0; extra == "dev"
37
+ Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
38
+ Requires-Dist: mypy>=1.0.0; extra == "dev"
39
+ Requires-Dist: build>=1.0.0; extra == "dev"
40
+ Requires-Dist: twine>=5.0.0; extra == "dev"
41
+
42
+ # Android Security Environment (`secenvY`)
43
+
44
+ A high-performance Gymnasium environment and Reinforcement Learning brain layer for AI-driven Android penetration testing, automated vulnerability discovery, and dynamic security analysis.
45
+
46
+ ---
47
+
48
+ ## 🌟 Key Features
49
+
50
+ - **Standard Gymnasium `Discrete(8)` Action Space:**
51
+ 1. `0: TAP` (UI Exploration) — Tap clickable element / primary interactable.
52
+ 2. `1: TYPE_TEXT` (UI Exploration) — Type contextual credentials / input fields.
53
+ 3. `2: SCROLL` (UI Exploration) — Scroll down to reveal off-screen views.
54
+ 4. `3: BACK` (UI Exploration) — Press Android system back button.
55
+ 5. `4: INJECT_SQL_PAYLOAD` (Security Testing) — Fuzz active input fields / queries with SQL syntax.
56
+ 6. `5: START_ACTIVITY` (Security Testing) — Launch target exported Activity component directly via `am start`.
57
+ 7. `6: QUERY_PROVIDER` (Security Testing) — Query target Content Provider URI directly.
58
+ 8. `7: FINISH` (Lifecycle) — Conclude security audit session cleanly with completion rewards.
59
+ - **117-Dimensional Normalized Observation Tensor:** Comprehensive 17-group state representation bounded strictly in `[0.0, 1.0]` (`gymnasium.spaces.Box(low=0.0, high=1.0, shape=(117,), dtype=np.float32)`).
60
+ - **Sub-Millisecond Pipeline (< 1ms SLA):** Highly optimized validation, building, and encoding pipeline executing end-to-end in ~0.89 ms.
61
+ - **Multi-Tier RAG-Driven Reward Engine:** Objective reward formulation combining step efficiency, screen discovery bonuses, severity-weighted RAG evaluations, and anti-farming deduplication.
62
+ - **Production WebSocket Gateway & Server:** Synchronous `WebSocketSystemGateway` and standalone `server.py` FastAPI routing hub bridging RL environments, physical Android devices, and desktop dashboards.
63
+
64
+ ---
65
+
66
+ ## 📦 Installation
67
+
68
+ ```bash
69
+ # Install from PyPI
70
+ pip install secenvY
71
+
72
+ # Or install from source in editable mode
73
+ git clone https://github.com/d4em0ny/android-security-agent.git
74
+ cd android-security-agent
75
+ python3 -m venv .venv
76
+ source .venv/bin/activate
77
+ pip install -e ".[dev]"
78
+ ```
79
+
80
+ ---
81
+
82
+ ## 🚀 Quickstart Usage
83
+
84
+ ```python
85
+ import gymnasium as gym
86
+ import secenvY
87
+
88
+ # 1. Create registered environment
89
+ env = gym.make("AndroidSecurityEnv-v0")
90
+
91
+ # 2. Reset environment to clean initial state
92
+ obs, info = env.reset(seed=42, options={"pkg": "com.vulnerable.bank"})
93
+ print(f"Observation Shape: {obs.shape}") # (117,)
94
+ print(f"Action Space: {env.action_space}") # Discrete(8)
95
+ print(f"Action Mask: {info['action_mask']}")
96
+
97
+ # 3. Step through episode
98
+ terminated = False
99
+ truncated = False
100
+ while not (terminated or truncated):
101
+ action_mask = info.get("action_mask")
102
+ # Sample or select discrete action (0..7)
103
+ action = env.action_space.sample(mask=action_mask)
104
+ obs, reward, terminated, truncated, info = env.step(action)
105
+ print(f"Step {info['step']} -> Reward: {reward:+.3f} (Terminated: {terminated})")
106
+
107
+ env.close()
108
+ ```
109
+
110
+ ---
111
+
112
+ ## 🧪 Running Tests & Verification
113
+
114
+ ```bash
115
+ # Run all unit and integration tests
116
+ pytest tests/ -v
117
+
118
+ # Run performance benchmark (< 5ms SLA verification)
119
+ pytest tests/test_pipeline_performance.py -v -s
120
+
121
+ # Run static type checker
122
+ mypy secenvY/ server.py main.py --config-file mypy.ini
123
+ ```
124
+
125
+ ---
126
+
127
+ ## 📄 License
128
+
129
+ Apache-2.0 License.
@@ -0,0 +1,88 @@
1
+ # Android Security Environment (`secenvY`)
2
+
3
+ A high-performance Gymnasium environment and Reinforcement Learning brain layer for AI-driven Android penetration testing, automated vulnerability discovery, and dynamic security analysis.
4
+
5
+ ---
6
+
7
+ ## 🌟 Key Features
8
+
9
+ - **Standard Gymnasium `Discrete(8)` Action Space:**
10
+ 1. `0: TAP` (UI Exploration) — Tap clickable element / primary interactable.
11
+ 2. `1: TYPE_TEXT` (UI Exploration) — Type contextual credentials / input fields.
12
+ 3. `2: SCROLL` (UI Exploration) — Scroll down to reveal off-screen views.
13
+ 4. `3: BACK` (UI Exploration) — Press Android system back button.
14
+ 5. `4: INJECT_SQL_PAYLOAD` (Security Testing) — Fuzz active input fields / queries with SQL syntax.
15
+ 6. `5: START_ACTIVITY` (Security Testing) — Launch target exported Activity component directly via `am start`.
16
+ 7. `6: QUERY_PROVIDER` (Security Testing) — Query target Content Provider URI directly.
17
+ 8. `7: FINISH` (Lifecycle) — Conclude security audit session cleanly with completion rewards.
18
+ - **117-Dimensional Normalized Observation Tensor:** Comprehensive 17-group state representation bounded strictly in `[0.0, 1.0]` (`gymnasium.spaces.Box(low=0.0, high=1.0, shape=(117,), dtype=np.float32)`).
19
+ - **Sub-Millisecond Pipeline (< 1ms SLA):** Highly optimized validation, building, and encoding pipeline executing end-to-end in ~0.89 ms.
20
+ - **Multi-Tier RAG-Driven Reward Engine:** Objective reward formulation combining step efficiency, screen discovery bonuses, severity-weighted RAG evaluations, and anti-farming deduplication.
21
+ - **Production WebSocket Gateway & Server:** Synchronous `WebSocketSystemGateway` and standalone `server.py` FastAPI routing hub bridging RL environments, physical Android devices, and desktop dashboards.
22
+
23
+ ---
24
+
25
+ ## 📦 Installation
26
+
27
+ ```bash
28
+ # Install from PyPI
29
+ pip install secenvY
30
+
31
+ # Or install from source in editable mode
32
+ git clone https://github.com/d4em0ny/android-security-agent.git
33
+ cd android-security-agent
34
+ python3 -m venv .venv
35
+ source .venv/bin/activate
36
+ pip install -e ".[dev]"
37
+ ```
38
+
39
+ ---
40
+
41
+ ## 🚀 Quickstart Usage
42
+
43
+ ```python
44
+ import gymnasium as gym
45
+ import secenvY
46
+
47
+ # 1. Create registered environment
48
+ env = gym.make("AndroidSecurityEnv-v0")
49
+
50
+ # 2. Reset environment to clean initial state
51
+ obs, info = env.reset(seed=42, options={"pkg": "com.vulnerable.bank"})
52
+ print(f"Observation Shape: {obs.shape}") # (117,)
53
+ print(f"Action Space: {env.action_space}") # Discrete(8)
54
+ print(f"Action Mask: {info['action_mask']}")
55
+
56
+ # 3. Step through episode
57
+ terminated = False
58
+ truncated = False
59
+ while not (terminated or truncated):
60
+ action_mask = info.get("action_mask")
61
+ # Sample or select discrete action (0..7)
62
+ action = env.action_space.sample(mask=action_mask)
63
+ obs, reward, terminated, truncated, info = env.step(action)
64
+ print(f"Step {info['step']} -> Reward: {reward:+.3f} (Terminated: {terminated})")
65
+
66
+ env.close()
67
+ ```
68
+
69
+ ---
70
+
71
+ ## 🧪 Running Tests & Verification
72
+
73
+ ```bash
74
+ # Run all unit and integration tests
75
+ pytest tests/ -v
76
+
77
+ # Run performance benchmark (< 5ms SLA verification)
78
+ pytest tests/test_pipeline_performance.py -v -s
79
+
80
+ # Run static type checker
81
+ mypy secenvY/ server.py main.py --config-file mypy.ini
82
+ ```
83
+
84
+ ---
85
+
86
+ ## 📄 License
87
+
88
+ Apache-2.0 License.
@@ -0,0 +1,77 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "secenvY"
7
+ version = "0.1.0"
8
+ description = "Gymnasium Environment for AI-driven Android Security Analysis, Dynamic Fuzzing, and Penetration Testing"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = { text = "Apache-2.0" }
12
+ authors = [
13
+ { name = "Android Security Agent Team" },
14
+ ]
15
+ keywords = [
16
+ "reinforcement-learning",
17
+ "gymnasium",
18
+ "android",
19
+ "security",
20
+ "penetration-testing",
21
+ "fuzzing",
22
+ "rag",
23
+ "dueling-dqn",
24
+ ]
25
+ classifiers = [
26
+ "Development Status :: 4 - Beta",
27
+ "Intended Audience :: Developers",
28
+ "Intended Audience :: Science/Research",
29
+ "Topic :: Security",
30
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
31
+ "Programming Language :: Python :: 3",
32
+ "Programming Language :: Python :: 3.10",
33
+ "Programming Language :: Python :: 3.11",
34
+ "Programming Language :: Python :: 3.12",
35
+ ]
36
+ dependencies = [
37
+ "gymnasium>=0.29.0",
38
+ "numpy>=1.24.0",
39
+ "pydantic>=2.0.0",
40
+ "pydantic-settings>=2.0.0",
41
+ "torch>=2.0.0",
42
+ "chromadb>=0.5.0",
43
+ "google-genai>=0.1.0",
44
+ "sentence-transformers>=2.2.0",
45
+ "python-dotenv>=1.0.0",
46
+ "rich>=13.0.0",
47
+ "fastapi>=0.100.0",
48
+ "uvicorn>=0.20.0",
49
+ "websockets>=11.0",
50
+ ]
51
+
52
+ [project.optional-dependencies]
53
+ dev = [
54
+ "pytest>=7.0.0",
55
+ "pytest-cov>=4.0.0",
56
+ "mypy>=1.0.0",
57
+ "build>=1.0.0",
58
+ "twine>=5.0.0",
59
+ ]
60
+
61
+ [project.urls]
62
+ Homepage = "https://github.com/d4em0ny/android-security-agent"
63
+ Repository = "https://github.com/d4em0ny/android-security-agent.git"
64
+ Issues = "https://github.com/d4em0ny/android-security-agent/issues"
65
+
66
+ [tool.setuptools.packages.find]
67
+ where = ["."]
68
+ include = ["secenvY*"]
69
+
70
+ [tool.setuptools.package-data]
71
+ "secenvY" = ["data/**/*"]
72
+
73
+ [tool.pytest.ini_options]
74
+ testpaths = ["tests"]
75
+ python_files = ["test_*.py"]
76
+ python_functions = ["test_*"]
77
+ addopts = "-v --strict-markers"
@@ -0,0 +1,43 @@
1
+ """Android Security Analyzer Gymnasium Environment & Agent (SecEnv)."""
2
+
3
+ import gymnasium as gym
4
+
5
+ # Register standard AndroidSecurityEnv-v0 environment
6
+ gym.register(
7
+ id="AndroidSecurityEnv-v0",
8
+ entry_point="secenvY.env.security_env:SecurityEnv",
9
+ max_episode_steps=100,
10
+ )
11
+
12
+ from secenvY.agents.dqn_agent import DuelingDQNAgent
13
+ from secenvY.agents.networks import DuelingQNetwork
14
+ from secenvY.agents.replay_buffer import PrioritizedReplayBuffer, UniformReplayBuffer
15
+ from secenvY.env.security_env import SecurityEnv
16
+ from secenvY.gateway.base import SystemGateway
17
+ from secenvY.gateway.mock_gateway import MockSystemGateway
18
+ from secenvY.gateway.websocket_gateway import WebSocketSystemGateway
19
+ from secenvY.services.action_mapper import ActionMapper
20
+ from secenvY.services.episode_policy import EpisodePolicy
21
+ from secenvY.services.reward_calculator import RewardCalculator
22
+ from secenvY.services.state_builder import StateBuilder
23
+ from secenvY.services.state_encoder import StateEncoder
24
+ from secenvY.services.state_validator import StateValidator
25
+ from secenvY.services.tensor_converter import TensorConverter
26
+
27
+ __all__ = [
28
+ "SecurityEnv",
29
+ "DuelingDQNAgent",
30
+ "DuelingQNetwork",
31
+ "PrioritizedReplayBuffer",
32
+ "UniformReplayBuffer",
33
+ "SystemGateway",
34
+ "MockSystemGateway",
35
+ "WebSocketSystemGateway",
36
+ "ActionMapper",
37
+ "StateBuilder",
38
+ "StateEncoder",
39
+ "StateValidator",
40
+ "TensorConverter",
41
+ "RewardCalculator",
42
+ "EpisodePolicy",
43
+ ]
@@ -0,0 +1,17 @@
1
+ """Dueling Double Deep Q-Network agent, neural network architectures, and replay buffers."""
2
+
3
+ from secenvY.agents.dqn_agent import DuelingDQNAgent
4
+ from secenvY.agents.networks import DuelingQNetwork
5
+ from secenvY.agents.replay_buffer import (
6
+ PrioritizedReplayBuffer,
7
+ SumTree,
8
+ UniformReplayBuffer,
9
+ )
10
+
11
+ __all__ = [
12
+ "DuelingDQNAgent",
13
+ "DuelingQNetwork",
14
+ "PrioritizedReplayBuffer",
15
+ "UniformReplayBuffer",
16
+ "SumTree",
17
+ ]
@@ -0,0 +1,225 @@
1
+ """Dueling Double Deep Q-Network (Dueling Double DQN) Agent with Prioritized Experience Replay."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ import random
7
+ from typing import Any, Dict, List, Optional, Tuple, Union
8
+ import numpy as np
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.optim as optim
12
+
13
+ from secenvY.agents.networks import DuelingQNetwork
14
+ from secenvY.agents.replay_buffer import PrioritizedReplayBuffer, UniformReplayBuffer
15
+
16
+
17
+ class DuelingDQNAgent:
18
+ """Reinforcement Learning Agent implementing Dueling Double DQN with Prioritized Experience Replay (PER).
19
+
20
+ Args:
21
+ state_dim: Dimension of observation tensor (default: 117).
22
+ action_dim: Number of discrete actions (default: 8).
23
+ hidden_dim: Hidden dimension of backbone network (default: 256).
24
+ lr: Learning rate for Adam optimizer (default: 1e-4).
25
+ gamma: Discount factor for Bellman target (default: 0.99).
26
+ tau: Soft target network update coefficient (default: 0.005).
27
+ target_update_interval: Hard update interval if tau is 1.0 (default: 500).
28
+ buffer_capacity: Max transitions stored in experience replay (default: 100,000).
29
+ use_per: Whether to use Prioritized Experience Replay (default: True).
30
+ per_alpha: PER prioritization exponent (default: 0.6).
31
+ per_beta_start: Initial importance sampling exponent (default: 0.4).
32
+ max_grad_norm: Maximum gradient norm for clipping (default: 10.0).
33
+ device: Torch computing device (default: cpu / cuda if available).
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ state_dim: int = 117,
39
+ action_dim: int = 8,
40
+ hidden_dim: int = 256,
41
+ lr: float = 1e-4,
42
+ gamma: float = 0.99,
43
+ tau: float = 0.005,
44
+ target_update_interval: int = 500,
45
+ buffer_capacity: int = 100_000,
46
+ use_per: bool = True,
47
+ per_alpha: float = 0.6,
48
+ per_beta_start: float = 0.4,
49
+ max_grad_norm: float = 10.0,
50
+ device: Optional[Union[str, torch.device]] = None,
51
+ ) -> None:
52
+ if device is None:
53
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
54
+ else:
55
+ self.device = torch.device(device)
56
+
57
+ self.state_dim = state_dim
58
+ self.action_dim = action_dim
59
+ self.gamma = gamma
60
+ self.tau = tau
61
+ self.target_update_interval = target_update_interval
62
+ self.max_grad_norm = max_grad_norm
63
+ self.use_per = use_per
64
+
65
+ # Neural Networks: Online and Target Q-Networks
66
+ self.online_net = DuelingQNetwork(state_dim, action_dim, hidden_dim).to(self.device)
67
+ self.target_net = DuelingQNetwork(state_dim, action_dim, hidden_dim).to(self.device)
68
+ self.target_net.load_state_dict(self.online_net.state_dict())
69
+ self.target_net.eval()
70
+ for param in self.target_net.parameters():
71
+ param.requires_grad = False
72
+
73
+ # Optimizer and Huber Loss Criterion
74
+ self.optimizer = optim.Adam(self.online_net.parameters(), lr=lr, eps=1e-8)
75
+ self.criterion = nn.SmoothL1Loss(reduction="none")
76
+
77
+ # Replay Memory (PER or Uniform)
78
+ self.replay_buffer: Union[PrioritizedReplayBuffer, UniformReplayBuffer]
79
+ if use_per:
80
+ self.replay_buffer = PrioritizedReplayBuffer(
81
+ capacity=buffer_capacity,
82
+ alpha=per_alpha,
83
+ beta_start=per_beta_start,
84
+ device=self.device,
85
+ )
86
+ else:
87
+ self.replay_buffer = UniformReplayBuffer(
88
+ capacity=buffer_capacity,
89
+ device=self.device,
90
+ )
91
+
92
+ self.total_steps = 0
93
+ self.train_steps = 0
94
+
95
+ def select_action(
96
+ self,
97
+ state: np.ndarray | List[float],
98
+ epsilon: float = 0.0,
99
+ action_mask: Optional[Union[np.ndarray, List[bool]]] = None,
100
+ ) -> int:
101
+ """Select action using epsilon-greedy exploration policy with optional action masking.
102
+
103
+ Args:
104
+ state: State vector of shape (117,) or list of floats.
105
+ epsilon: Exploration probability in [0.0, 1.0].
106
+ action_mask: Optional boolean mask of shape (action_dim,) where True indicates valid action.
107
+
108
+ Returns:
109
+ Discrete action ID in range [0, action_dim - 1].
110
+ """
111
+ self.total_steps += 1
112
+
113
+ if action_mask is not None:
114
+ mask_arr = np.array(action_mask, dtype=bool)
115
+ valid_indices = np.where(mask_arr)[0]
116
+ if len(valid_indices) == 0:
117
+ valid_indices = np.arange(self.action_dim)
118
+ else:
119
+ valid_indices = np.arange(self.action_dim)
120
+
121
+ if random.random() < epsilon:
122
+ return int(random.choice(valid_indices))
123
+
124
+ state_tensor = torch.tensor(
125
+ np.array(state, dtype=np.float32),
126
+ device=self.device,
127
+ dtype=torch.float32,
128
+ ).unsqueeze(0)
129
+
130
+ self.online_net.eval()
131
+ with torch.no_grad():
132
+ q_values = self.online_net(state_tensor).squeeze(0)
133
+ if action_mask is not None:
134
+ mask_tensor = torch.tensor(mask_arr, device=self.device, dtype=torch.bool)
135
+ q_values[~mask_tensor] = -1e9
136
+ action = int(q_values.argmax(dim=-1).item())
137
+ self.online_net.train()
138
+
139
+ return action
140
+
141
+ def train_step(self, batch_size: int = 32) -> Optional[float]:
142
+ """Perform a single Double DQN training step over a batch from replay buffer.
143
+
144
+ Args:
145
+ batch_size: Number of transitions to sample.
146
+
147
+ Returns:
148
+ Computed scalar loss value, or None if buffer has insufficient transitions.
149
+ """
150
+ if len(self.replay_buffer) < batch_size:
151
+ return None
152
+
153
+ self.train_steps += 1
154
+
155
+ # 1. Sample batch
156
+ if self.use_per:
157
+ states, actions, rewards, next_states, dones, indices, weights = self.replay_buffer.sample(batch_size) # type: ignore
158
+ else:
159
+ states, actions, rewards, next_states, dones = self.replay_buffer.sample(batch_size) # type: ignore
160
+ weights = torch.ones_like(rewards)
161
+ indices = []
162
+
163
+ # 2. Compute current Q-values Q(s, a; θ)
164
+ current_q = self.online_net(states).gather(1, actions)
165
+
166
+ # 3. Double DQN Target: a* = argmax_{a'} Q(s', a'; θ), y = r + γ Q(s', a*; θ^-)
167
+ with torch.no_grad():
168
+ next_actions = self.online_net(next_states).argmax(dim=-1, keepdim=True)
169
+ target_q_next = self.target_net(next_states).gather(1, next_actions)
170
+ target_q = rewards + (1.0 - dones) * self.gamma * target_q_next
171
+
172
+ # 4. Compute TD Errors and Huber Loss
173
+ td_errors = (current_q - target_q).detach()
174
+ loss_elements = self.criterion(current_q, target_q)
175
+ loss = (loss_elements * weights).mean()
176
+
177
+ # 5. Backpropagation and Gradient Clipping
178
+ self.optimizer.zero_grad()
179
+ loss.backward()
180
+ if self.max_grad_norm > 0:
181
+ nn.utils.clip_grad_norm_(self.online_net.parameters(), max_norm=self.max_grad_norm)
182
+ self.optimizer.step()
183
+
184
+ # 6. Update PER priorities
185
+ if self.use_per and indices:
186
+ self.replay_buffer.update_priorities(indices, td_errors.abs()) # type: ignore
187
+
188
+ # 7. Update Target Network
189
+ self._update_target_network()
190
+
191
+ return float(loss.item())
192
+
193
+ def _update_target_network(self) -> None:
194
+ """Update target network parameters via Polyak soft update or periodic hard copy."""
195
+ if self.tau < 1.0:
196
+ # Soft Polyak update: θ^- ← τ θ + (1 - τ) θ^-
197
+ for target_param, online_param in zip(self.target_net.parameters(), self.online_net.parameters()):
198
+ target_param.data.copy_(self.tau * online_param.data + (1.0 - self.tau) * target_param.data)
199
+ else:
200
+ # Hard periodic update
201
+ if self.train_steps % self.target_update_interval == 0:
202
+ self.target_net.load_state_dict(self.online_net.state_dict())
203
+
204
+ def save_checkpoint(self, filepath: str) -> None:
205
+ """Save model and optimizer state to disk."""
206
+ os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
207
+ torch.save(
208
+ {
209
+ "online_net_state": self.online_net.state_dict(),
210
+ "target_net_state": self.target_net.state_dict(),
211
+ "optimizer_state": self.optimizer.state_dict(),
212
+ "total_steps": self.total_steps,
213
+ "train_steps": self.train_steps,
214
+ },
215
+ filepath,
216
+ )
217
+
218
+ def load_checkpoint(self, filepath: str) -> None:
219
+ """Load model and optimizer state from disk."""
220
+ checkpoint = torch.load(filepath, map_location=self.device)
221
+ self.online_net.load_state_dict(checkpoint["online_net_state"])
222
+ self.target_net.load_state_dict(checkpoint["target_net_state"])
223
+ self.optimizer.load_state_dict(checkpoint["optimizer_state"])
224
+ self.total_steps = checkpoint.get("total_steps", 0)
225
+ self.train_steps = checkpoint.get("train_steps", 0)
@@ -0,0 +1,81 @@
1
+ """PyTorch neural network architectures for Dueling Deep Q-Networks."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+
8
+
9
+ class DuelingQNetwork(nn.Module):
10
+ """Dueling Deep Q-Network decomposing Q(s, a) into Value V(s) and Advantage A(s, a) streams.
11
+
12
+ Q(s, a) = V(s) + (A(s, a) - 1/|A| * sum_{a'} A(s, a'))
13
+
14
+ Args:
15
+ state_dim: Input observation space dimension (default: 117).
16
+ action_dim: Discrete action space size (default: 8).
17
+ hidden_dim: Number of hidden units in backbone layers (default: 256).
18
+ """
19
+
20
+ def __init__(
21
+ self,
22
+ state_dim: int = 117,
23
+ action_dim: int = 8,
24
+ hidden_dim: int = 256,
25
+ ) -> None:
26
+ super().__init__()
27
+ self.state_dim = state_dim
28
+ self.action_dim = action_dim
29
+
30
+ # Shared Feature Extractor Backbone
31
+ self.backbone = nn.Sequential(
32
+ nn.Linear(state_dim, hidden_dim),
33
+ nn.LayerNorm(hidden_dim),
34
+ nn.ReLU(),
35
+ nn.Linear(hidden_dim, hidden_dim),
36
+ nn.ReLU(),
37
+ )
38
+
39
+ # Value Stream: V(s) -> [B, 1]
40
+ self.value_stream = nn.Sequential(
41
+ nn.Linear(hidden_dim, hidden_dim // 2),
42
+ nn.ReLU(),
43
+ nn.Linear(hidden_dim // 2, 1),
44
+ )
45
+
46
+ # Advantage Stream: A(s, a) -> [B, action_dim]
47
+ self.advantage_stream = nn.Sequential(
48
+ nn.Linear(hidden_dim, hidden_dim // 2),
49
+ nn.ReLU(),
50
+ nn.Linear(hidden_dim // 2, action_dim),
51
+ )
52
+
53
+ def forward(self, state: torch.Tensor) -> torch.Tensor:
54
+ """Forward pass computing Q-values for all discrete actions.
55
+
56
+ Args:
57
+ state: Tensor of shape (batch_size, state_dim) or (state_dim,).
58
+
59
+ Returns:
60
+ Tensor of shape (batch_size, action_dim) containing Q-values.
61
+ """
62
+ if state.dim() == 1:
63
+ state = state.unsqueeze(0)
64
+
65
+ features = self.backbone(state)
66
+ value = self.value_stream(features) # [B, 1]
67
+ advantage = self.advantage_stream(features) # [B, action_dim]
68
+
69
+ # Aggregation: Q(s, a) = V(s) + (A(s, a) - mean(A(s, :)))
70
+ q_values = value + (advantage - advantage.mean(dim=-1, keepdim=True))
71
+ return q_values
72
+
73
+ def compute_value_and_advantage(self, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
74
+ """Expose raw value and advantage streams for inspection and explainability."""
75
+ if state.dim() == 1:
76
+ state = state.unsqueeze(0)
77
+
78
+ features = self.backbone(state)
79
+ value = self.value_stream(features)
80
+ advantage = self.advantage_stream(features)
81
+ return value, advantage