File size: 6,445 Bytes
323c6ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | # LINE 1-17: ARE USED TO SILENCE LOGS & ALIGN NAMESPACES
import os
import warnings
import tensorflow as tf
os.environ["TF_USE_LEGACY_KERAS"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["TF_FUNCTION_NUMERIC_CHECKS"] = "0"
tf.keras.backend.clear_session()
tf.config.threading.set_intra_op_parallelism_threads(0)
tf.config.threading.set_inter_op_parallelism_threads(0)
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
warnings.filterwarnings("ignore", category=UserWarning, module="gym")
warnings.filterwarnings("ignore", category=RuntimeWarning, module="tf_agents.specs.array_spec")
warnings.filterwarnings("ignore", category=DeprecationWarning)
from tf_agents.agents.ppo import ppo_agent
from tf_agents.networks import network
from tf_agents.networks import categorical_projection_network
from tf_agents.environments import tf_py_environment
# Custom Modular Imports
from config import Config
from env import CyberPettingZooEnv
from utilities import configure_gpu, PettingZooToTFAgentsWrapper
from trainer import MultiAgentTrainer
# BULLETPROOF CUSTOM GNN POLICY NETWORK FOR TF-AGENTS
class CyberGNNNetwork(network.Network):
"""
A custom TF-Agents Network that binds GNN feature extraction with
native CategoricalProjectionNetworks, explicitly forwarding outer_rank.
"""
def __init__(self, input_tensor_spec, output_spec, fc_layer_params=(128, 64), is_value_net=False, name="CyberGNNNetwork"):
super().__init__(input_tensor_spec=input_tensor_spec, state_spec=(), name=name)
self.is_value_net = is_value_net
# 1. Instantiate the Graph Convolution Weights
self.w1 = tf.Variable(tf.keras.initializers.GlorotUniform()(shape=(2, 32)), trainable=True, name="gcn_w1")
self.w2 = tf.Variable(tf.keras.initializers.GlorotUniform()(shape=(32, 32)), trainable=True, name="gcn_w2")
# 2. Instantiate standard Dense MLPs for decision making
self.dense_layers = []
for units in fc_layer_params:
self.dense_layers.append(tf.keras.layers.Dense(units, activation=tf.nn.tanh))
# 3. Instantiate native TF-Agents projection heads
if self.is_value_net:
self.projection_head = tf.keras.layers.Dense(1, kernel_initializer=tf.keras.initializers.Orthogonal(1.0))
else:
self.projection_head = categorical_projection_network.CategoricalProjectionNetwork(output_spec)
def call(self, observations, step_type=(), network_state=(), training=False):
A_norm = observations['adjacency_matrix']
X = observations['node_features']
outer_rank = X.shape.rank - 2
# GNN Message-Passing Block
h1 = tf.nn.tanh(tf.matmul(tf.matmul(A_norm, X), self.w1))
h2 = tf.nn.tanh(tf.matmul(tf.matmul(A_norm, h1), self.w2))
# Dynamically calculate batch structure to prevent rank mismatch crashes
batch_shape = tf.shape(h2)[:-2]
flat_features = tf.reshape(h2, tf.concat([batch_shape, [-1]], axis=0))
# Standard Policy Dense Layer Block
x = flat_features
for layer in self.dense_layers:
x = layer(x)
# Output Generation Pass
if self.is_value_net:
value_predictions = self.projection_head(x)
value_predictions = tf.squeeze(value_predictions, axis=-1)
return value_predictions, network_state
else:
action_distributions, _ = self.projection_head(x, outer_rank=outer_rank, training=training)
return action_distributions, network_state
# AGENT BUILD ORCHESTRATION WITH EXPLICIT VARIABLE REGISTRATION
def create_ppo_agent(env_time_step_spec, env_action_spec, agent_name, lr_schedule):
"""Creates an independent PPO Agent with guaranteed trainable variable tracking."""
agent_obs_spec = env_time_step_spec.observation[agent_name]
agent_action_spec = env_action_spec[agent_name]
actor_net = CyberGNNNetwork(
input_tensor_spec=agent_obs_spec,
output_spec=agent_action_spec,
fc_layer_params=Config.ACTOR_LAYERS,
is_value_net=False,
name=f"actor_gnn_{agent_name}"
)
value_net = CyberGNNNetwork(
input_tensor_spec=agent_obs_spec,
output_spec=None,
fc_layer_params=Config.CRITIC_LAYERS,
is_value_net=True,
name=f"value_gnn_{agent_name}"
)
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
agent = ppo_agent.PPOAgent(
time_step_spec=env_time_step_spec._replace(
observation=agent_obs_spec,
reward=env_time_step_spec.reward[agent_name]
),
action_spec=agent_action_spec,
optimizer=optimizer,
actor_net=actor_net,
value_net=value_net,
num_epochs=10,
discount_factor=Config.GAMMA,
use_gae=True,
use_td_lambda_return=True,
normalize_observations=False,
normalize_rewards=True,
value_pred_loss_coef=0.5
)
agent.initialize()
print(f"🔒 Verified Trainable Variables for {agent_name} Actor: {len(actor_net.trainable_variables)}")
return agent
# MAIN RUNTIME EXECUTION LOOP
def main():
configure_gpu()
print(f"🚀 Initializing Custom GNN MARL Cyber Range Framework in **{Config.MODE}** mode.")
raw_pz_env = CyberPettingZooEnv()
py_env = PettingZooToTFAgentsWrapper(raw_pz_env)
tf_env = tf_py_environment.TFPyEnvironment(py_env)
time_step_spec = tf_env.time_step_spec()
action_spec = tf_env.action_spec()
print("🧠 Constructing Custom Graph Policy Networks for Red Agent...")
red_agent = create_ppo_agent(time_step_spec, action_spec, 'red', Config.get_red_lr_schedule())
print("🧠 Constructing Custom Graph Policy Networks for Blue Agent...")
blue_agent = create_ppo_agent(time_step_spec, action_spec, 'blue', Config.get_blue_lr_schedule())
trainer = MultiAgentTrainer(config=Config, tf_env=tf_env, agent_1=red_agent, agent_2=blue_agent)
print("✨ Graph Neural Network pipeline successfully compiled. Launching training loop...")
for epoch in range(1, Config.TOTAL_EPISODES + 1):
trainer.train_epoch(epoch_idx=epoch)
if __name__ == '__main__':
main() |