CyberBattleSim `agents` 模块详细分析

CyberBattleSim agents 模块详细分析

agents 模块是 CyberBattleSim 的核心组件之一，包含了各种智能体实现。下面我将详细分析每个文件和类。

1. agents/init.py

这个文件将模块中的关键类暴露给外部使用。

导出的类

• Learner：学习器基类
• AgentWrapper：智能体包装器
• EnvironmentBounds：环境边界参数

2. agents/baseline/agent_wrapper.py

这个文件定义了各种特征提取器和状态增强类，用于将环境观察转换为智能体可用的特征。

StateAugmentation 类

• 主要作用：基础状态增强类，存储和更新环境观察
• 初始化参数：
  • observation：初始环境观察
• 主要方法：
  • on_reset(observation)：重置状态
  • on_step(action, reward, done, truncated, observation)：更新状态

Feature 类

• 主要作用：特征提取器的抽象基类
• 初始化参数：
  • p：环境边界参数
  • nvec：特征向量维度
• 主要方法：
  • get(a)：获取特征向量
  • get_at(a, node)：获取特定节点的特征向量
  • name()：获取特征名称

各种特征提取器类

这些类都继承自 Feature，用于提取不同类型的特征：

Feature_discovered_node_count

• 主要作用：提取已发现节点数量
• 初始化参数：p：环境边界参数
• 处理逻辑：返回已发现节点数量的特征向量

Feature_active_node_age

• 主要作用：提取当前节点的发现年龄
• 初始化参数：p：环境边界参数
• 处理逻辑：返回节点在发现列表中的位置

Feature_active_node_properties

• 主要作用：提取当前节点的属性
• 初始化参数：p：环境边界参数
• 处理逻辑：返回节点属性的二进制向量

Feature_discovered_ports

• 主要作用：提取已发现的端口
• 初始化参数：p：环境边界参数
• 处理逻辑：返回已发现端口的二进制向量

Feature_discovered_credential_count

• 主要作用：提取已发现的凭证数量
• 初始化参数：p：环境边界参数
• 处理逻辑：返回已发现凭证数量的特征向量

Feature_owned_node_count

• 主要作用：提取已拥有节点数量
• 初始化参数：p：环境边界参数
• 处理逻辑：返回已拥有节点数量的特征向量

Feature_discovered_notowned_node_count

• 主要作用：提取已发现但未拥有的节点数量
• 初始化参数：
  • p：环境边界参数
  • node：目标节点（可选）
• 处理逻辑：返回已发现但未拥有节点数量的特征向量

ConcatFeatures

• 主要作用：将多个特征连接成一个特征向量
• 初始化参数：
  • p：环境边界参数
  • feature_selection：要连接的特征列表
• 处理逻辑：连接多个特征向量，返回一个大的特征向量

HashFeatures

• 主要作用：将特征向量哈希到固定大小
• 初始化参数：
  • p：环境边界参数
  • feature_selection：要哈希的特征列表
  • hash_size：哈希空间大小
• 处理逻辑：将特征向量哈希到指定大小的空间

RavelFeatures

• 主要作用：将多维特征向量展平
• 初始化参数：
  • p：环境边界参数
  • feature_selection：要展平的特征列表
• 处理逻辑：将多维特征向量展平为一维向量

AbstractAction 类

• 主要作用：抽象动作空间，简化动作表示
• 初始化参数：p：环境边界参数
• 主要方法：
  • get_action(abstract_action, source_node)：获取具体动作
  • get_abstract_action(action)：获取抽象动作
• 处理逻辑：将复杂的动作空间抽象为简单的索引

ActionTrackingStateAugmentation 类

• 主要作用：跟踪动作成功和失败的状态增强
• 初始化参数：
  • p：环境边界参数
  • observation：初始环境观察
• 主要方法：
  • on_step(action, reward, done, truncated, observation)：更新动作跟踪状态
• 处理逻辑：记录每个节点上动作的成功和失败次数

Feature_success_actions_at_node 类

• 主要作用：提取节点上成功动作的特征
• 初始化参数：p：环境边界参数
• 处理逻辑：返回节点上成功动作的二进制向量

Feature_failed_actions_at_node 类

• 主要作用：提取节点上失败动作的特征
• 初始化参数：p：环境边界参数
• 处理逻辑：返回节点上失败动作的二进制向量

Feature_actions_tried_at_node 类

• 主要作用：提取节点上尝试过的动作特征
• 初始化参数：p：环境边界参数
• 处理逻辑：返回节点上尝试过的动作的二进制向量

AgentWrapper 类

• 主要作用：包装环境，更新智能体状态
• 初始化参数：
  • env：要包装的环境
  • state：状态增强对象
• 主要方法：
  • step(action)：执行动作并更新状态
  • reset(**kwargs)：重置环境和状态
• 处理逻辑：在每次环境交互后更新智能体状态

3. agents/baseline/agent_dql.py

这个文件实现了深度 Q 学习智能体。

Transition 类

• 主要作用：存储经验回放的转换
• 处理逻辑：作为 NamedTuple，存储状态、动作、下一状态和奖励

ReplayMemory 类

• 主要作用：经验回放缓冲区
• 初始化参数：capacity：缓冲区容量
• 主要方法：
  • push(*args)：添加转换
  • sample(batch_size)：随机采样批次
  • __len__()：获取当前大小
• 处理逻辑：存储和采样转换，实现经验回放

CyberBattleStateActionModel 类

• 主要作用：定义状态和动作空间模型
• 初始化参数：ep：环境边界参数
• 主要方法：
  • get_actor_state_vector(global_state, actor_features)：获取状态向量
  • get_state_vector(a, node)：获取特定节点的状态向量
  • get_state_action_vector(a, node, abstract_action)：获取状态-动作向量
• 处理逻辑：提取和组合状态和动作特征

DQN 类

• 主要作用：深度 Q 网络模型
• 初始化参数：ep：环境边界参数
• 主要方法：
  • forward(x)：前向传播
• 处理逻辑：实现神经网络，计算 Q 值

CredentialCacheExploiter 类

• 主要作用：利用凭证缓存的策略
• 主要方法：
  • exploit(env, observation)：利用凭证缓存
• 处理逻辑：尝试使用已发现的凭证连接到新节点

DeepQLearnerPolicy 类

• 主要作用：深度 Q 学习策略
• 初始化参数：
  • ep：环境边界参数
  • gamma：折扣因子
  • replay_memory_size：回放缓冲区大小
  • target_update：目标网络更新频率
  • batch_size：批次大小
  • learning_rate：学习率
• 主要方法：
  • parameters_as_string()：获取参数字符串
  • all_parameters_as_string()：获取所有参数字符串
  • exploit(env, observation)：利用策略
  • learn(env, observation, action, reward, next_observation, done)：学习
  • optimize_model()：优化模型
• 处理逻辑：实现深度 Q 学习算法，包括经验回放、目标网络和优化

4. agents/baseline/agent_randomcredlookup.py

RandomCredentialCacheLookup 类

• 主要作用：随机凭证缓存查找策略
• 初始化参数：无
• 主要方法：
  • exploit(env, observation)：利用策略
• 处理逻辑：随机选择已发现的凭证尝试连接到新节点

5. agents/baseline/agent_tabularqlearning.py

TabularQLearnerPolicy 类

• 主要作用：表格 Q 学习策略
• 初始化参数：
  • ep：环境边界参数
  • learning_rate：学习率
  • gamma：折扣因子
  • epsilon：探索率
  • epsilon_decay：探索率衰减
  • epsilon_min：最小探索率
  • feature_selection：特征选择
  • hash_size：哈希空间大小
• 主要方法：
  • parameters_as_string()：获取参数字符串
  • exploit(env, observation)：利用策略
  • learn(env, observation, action, reward, next_observation, done)：学习
• 处理逻辑：实现表格 Q 学习算法，包括状态抽象和 Q 表更新

ChosenActionMetadata 类

• 主要作用：存储选择动作的元数据
• 处理逻辑：作为 NamedTuple，存储 Q 值和状态信息

6. agents/baseline/learner.py

Learner 类

• 主要作用：学习器抽象基类
• 主要方法：
  • exploit(env, observation)：利用策略
  • learn(env, observation, action, reward, next_observation, done)：学习
• 处理逻辑：定义学习器接口

epsilon_greedy_search 函数

• 主要作用：实现 epsilon-贪婪搜索
• 参数：
  • cyberbattle_gym_env：环境
  • environment_properties：环境属性
  • learner：学习器
  • episode_count：回合数
  • iteration_count：迭代数
  • epsilon：探索率
  • epsilon_exponential_decay：探索率指数衰减
  • epsilon_minimum：最小探索率
  • verbosity：详细程度
  • render：是否渲染
  • plot_episodes_length：是否绘制回合长度
  • title：标题
• 处理逻辑：实现 epsilon-贪婪策略，平衡探索和利用

7. agents/random_agent.py

RandomAgent 类

• 主要作用：随机动作智能体
• 初始化参数：无
• 主要方法：
  • run(env, render=False, episodes=1, iterations=100)：运行智能体
• 处理逻辑：随机选择有效动作

总结

CyberBattleSim 的 agents 模块提供了多种智能体实现，从简单的随机智能体到复杂的深度 Q 学习智能体。这些智能体通过不同的策略和学习算法，在网络安全模拟环境中进行决策和学习。模块的设计非常模块化，通过特征提取、状态增强和动作抽象等组件，使得智能体能够处理复杂的网络安全环境。

CyberBattleSim 中 agents 文件夹的文件和类关联分析

agents 文件夹包含了 CyberBattleSim 中所有智能体的实现，这些文件和类之间有着紧密的关联。下面我将详细分析这些关联关系。

文件之间的关联

1. 模块层次结构

• agents/__init__.py: 导出主要类，使它们可以直接从 cyberbattle.agents 导入
• agents/random_agent.py: 实现简单的随机智能体
• agents/baseline/: 包含更复杂的基准智能体实现
  • baseline/__init__.py: 导出基准模块的主要类
  • baseline/learner.py: 定义学习器接口和通用学习算法
  • baseline/agent_wrapper.py: 提供环境包装和特征提取
  • baseline/agent_dql.py: 实现深度 Q 学习智能体
  • baseline/agent_tabularqlearning.py: 实现表格 Q 学习智能体
  • baseline/agent_randomcredlookup.py: 实现随机凭证查找智能体

2. 依赖关系

• learner.py 是基础，定义了学习器接口
• agent_wrapper.py 提供特征提取和状态表示，被所有学习型智能体使用
• agent_dql.py 和 agent_tabularqlearning.py 依赖于 learner.py 和 agent_wrapper.py
• agent_randomcredlookup.py 是一个简单的策略，只依赖于基本环境接口

类之间的关联

1. 继承关系

• Learner (learner.py) 是所有学习型智能体的基类

  • TabularQLearnerPolicy (agent_tabularqlearning.py) 继承自 Learner
  • DeepQLearnerPolicy (agent_dql.py) 继承自 Learner

• Feature (agent_wrapper.py) 是所有特征提取器的基类

  • 各种特征提取器（如 Feature_discovered_node_count）继承自 Feature
  • 组合特征类（如 ConcatFeatures）也继承自 Feature

• StateAugmentation (agent_wrapper.py) 是状态增强的基类

  • ActionTrackingStateAugmentation 继承自 StateAugmentation

2. 组合关系

• DeepQLearnerPolicy 使用 ReplayMemory 存储经验
• DeepQLearnerPolicy 使用 DQN 作为神经网络模型
• TabularQLearnerPolicy 使用 HashFeatures 进行状态抽象
• AgentWrapper 包含 StateAugmentation 对象

3. 使用关系

• 所有学习型智能体使用 AbstractAction 简化动作空间
• epsilon_greedy_search 函数使用 Learner 接口进行训练
• 所有智能体通过 AgentWrapper 与环境交互

具体例子：深度 Q 学习智能体的工作流程

以下是一个具体例子，展示了这些类如何协同工作：

1. 在 dql_active_directory.py 笔记本中，创建了一个 DeepQLearnerPolicy 实例：

_l = dqla.DeepQLearnerPolicy(
    ep=ep,
    gamma=0.015,
    replay_memory_size=10000,
    target_update=5,
    batch_size=512,
    learning_rate=0.01,
)

2. 然后使用 epsilon_greedy_search 函数训练这个智能体：

dqn_learning_run = learner.epsilon_greedy_search(
    cyberbattle_gym_env=env,
    environment_properties=ep,
    learner=_l,
    episode_count=training_episode_count,
    iteration_count=iteration_count,
    epsilon=epsilon,
    epsilon_exponential_decay=50000,
    epsilon_minimum=0.1,
    verbosity=Verbosity.Quiet,
    render=False,
    plot_episodes_length=False,
    title=f"DQL {i}",
)

3. 训练过程中，epsilon_greedy_search 函数会：

   • 使用 AgentWrapper 包装环境，增强状态表示
   • 在每一步，以 ε 概率随机选择动作，或调用智能体的 exploit 方法
   • 执行选择的动作，获取新的观察和奖励
   • 调用智能体的 learn 方法更新模型

4. 在 DeepQLearnerPolicy.learn 方法中：

   • 使用 CyberBattleStateActionModel 提取状态和动作特征
   • 将经验存储到 ReplayMemory 中
   • 调用 optimize_model 方法更新 Q 网络

5. 在 DeepQLearnerPolicy.exploit 方法中：

   • 使用 CyberBattleStateActionModel 提取当前状态特征
   • 使用 Q 网络预测每个动作的 Q 值
   • 选择 Q 值最高的动作

这个例子展示了如何使用 DeepQLearnerPolicy 智能体在 Active Directory 环境中进行训练和测试。整个过程涉及多个类的协作，包括特征提取、状态表示、经验回放、神经网络模型等。

通过这种模块化设计，CyberBattleSim 提供了一个灵活的框架，可以轻松实现和测试不同的智能体算法。