SARSA.py

import gym
import numpy as np
from collections import deque

class NetworkEnv(gym.Env):
    def __init__(self, num_agents, num_actions):
        self.num_agents = num_agents
        self.num_actions = num_actions
        
        self.action_space = [gym.spaces.Discrete(num_actions) for _ in range(num_agents)]
        self.observation_space = [gym.spaces.Box(low=0, high=1, shape=(num_agents + 2,)) for _ in range(num_agents)]
        
        self.traffic_rate = 0
        self.attack_pattern = 0
        self.server_responsiveness = 1
        
    def step(self, actions):
        states = []
        rewards = []
        for agent_id, action in enumerate(actions):
            # Update the environment based on the agents' actions
            self.update_environment(agent_id, action)
            
            # Calculate the reward for the agent
            reward = self.calculate_reward(agent_id)
            rewards.append(reward)
            
            state = self.update_state(agent_id)
            states.append(state)
        
        return states, rewards, False, {}
    
    def reset(self):
        self.traffic_rate = 0
        self.attack_pattern = 0
        self.server_responsiveness = 1
        
        return [np.zeros(self.num_agents + 2) for _ in range(self.num_agents)]
    
    def calculate_reward(self, agent_id):
        return self.server_responsiveness - self.traffic_rate - self.attack_pattern
    
    def update_state(self, agent_id):
        # Update the state for the agent based on the current environment and the actions of other agents
        return np.concatenate((np.array([self.traffic_rate, self.attack_pattern]), np.random.rand(self.num_agents)))

class SarsaAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        self.q_table = np.zeros((state_size, action_size))
        self.gamma = 0.95
        self.alpha = 0.1
        
    def act(self, state):
        if np.random.rand() < 0.1:
            return np.random.randint(0, self.action_size)
        else:
            return np.argmax(self.q_table[state])
    
    def learn(self, state, action, reward, next_state, next_action):
        # Update the Q-table using the SARSA algorithm
        self.q_table[state, action] += self.alpha * (reward + self.gamma * self.q_table[next_state, next_action] - self.q_table[state, action])

env = NetworkEnv(num_agents=10, num_actions=5)
agents = [SarsaAgent(env.observation_space[i].shape[0], env.action_space[i].n) for i in range(env.num_agents)]

for episode in range(1000):
    states = env.reset()
    actions = [agent.act(state) for agent, state in zip(agents, states)]
    done = False
    while not done:
        next_states, rewards, done, _ = env.step(actions)
        next_actions = [agent.act(next_state) for agent, next_state in zip(agents, next_states)]
        for agent, state, action, reward, next_state, next_action in zip(agents, states, actions, rewards, next_states, next_actions):
            agent.learn(state, action, reward, next_state, next_action)
        states = next_states
        actions = next_actions