DDQN_CNN.py

import itertools
import math
import random
from collections import deque, namedtuple

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from simulation import *

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


class ReplayBuffer:

    def __init__(self, capacity):
        self.memory = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state):
        self.memory.append((state, action, reward, next_state))

    def sample(self, batch_size):
        batch = random.sample(self.memory, batch_size)

        states, actions, rewards, next_states = zip(*batch)

        states = torch.stack(states).float().to(device)
        actions = torch.tensor(actions, dtype=torch.long).to(device)
        rewards = torch.tensor(rewards, dtype=torch.float32).to(device)
        next_states = torch.stack(next_states).float().to(device)

        return states, actions, rewards, next_states

    def __len__(self):
        return len(self.memory)


class ActionSpace:

    def __init__(self, n_ues=2):
        self.n_ues = n_ues

        self.x_levels = [0, 1]  # offloading: local or remote
        self.beta_levels = [0.2, 0.5]  # computation allocation %
        self.theta_levels = [0.2, 0.5]  # spectrum allocation %

        # All possible choices for ONE UE (x, beta, theta)
        one_ue_choices = list(
            itertools.product(self.x_levels, self.beta_levels,
                              self.theta_levels))

        # Joint action list (tuples of per-UE triples)
        self.actions = list(itertools.product(one_ue_choices, repeat=n_ues))
        self.n_actions = len(self.actions)

    def decode(self, action_idx):
        """Return (x vector, beta vector, theta vector) for action index"""
        joint = self.actions[action_idx]
        xs, betas, thetas = [], [], []
        for (x, b, t) in joint:
            xs.append(x)
            betas.append(b)
            thetas.append(t)
        return np.array(xs), np.array(betas), np.array(thetas)


class CNN_DQN(nn.Module):

    def __init__(self, state_dim, n_actions):
        super(CNN_DQN, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv1d(in_channels=1, out_channels=32, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(in_channels=32, out_channels=64, kernel_size=1),
            nn.ReLU(),
        )

        # After conv layers, flatten
        self.fc = nn.Sequential(nn.Linear(64 * state_dim, 256), nn.ReLU(),
                                nn.Linear(256, 256), nn.ReLU(),
                                nn.Linear(256, n_actions))

    def forward(self, x):
        # x expected shape: (batch_size, state_dim)
        if x.dim() == 1:
            x = x.unsqueeze(0)  # → (1, state_dim)

        x = x.unsqueeze(1)  # → (batch, 1, state_dim)
        x = self.conv(x)  # → (batch, 64, state_dim)
        x = x.view(x.size(0), -1)  # flatten
        q = self.fc(x)  # → (batch, n_actions)
        return q


def run_save_and_plot_experiment(exp_name, n_runs=5):
    # Create folder for this experiment
    if not os.path.exists(exp_name):
        os.makedirs(exp_name)

    all_energies = []
    all_rewards = []
    all_losses = []

    for run in range(1, n_runs + 1):
        print('run: ', run)
        run_energies, run_rewards, run_losses = run_DQN()

        # Save CSV for this run
        df = pd.DataFrame({
            'Energy': run_energies,
            'Reward': run_rewards,
            'Loss': run_losses
        })
        df.to_csv(os.path.join(exp_name, f'run_{run}.csv'), index=False)

        # Collect for plotting
        all_energies.append(run_energies)
        all_rewards.append(run_rewards)
        all_losses.append(run_losses)

    all_energies = np.array(all_energies)
    all_rewards = np.array(all_rewards)
    all_losses = np.array(all_losses)

    episodes = np.arange(all_energies.shape[1])

    for data, name in zip([all_energies, all_rewards, all_losses],
                          ['Energy', 'Reward', 'Loss']):
        mean = data.mean(axis=0)
        std = data.std(axis=0)

        plt.figure()
        plt.plot(episodes, mean, label=f'Mean {name}')
        plt.fill_between(episodes,
                         mean - std,
                         mean + std,
                         alpha=0.3,
                         label='STD')
        plt.xlabel('Episode')
        plt.ylabel(name)
        plt.title(f'{name} Over Episodes ({exp_name})')
        plt.legend()
        plt.savefig(os.path.join(exp_name, f'{name}_plot.png'))
        plt.close()


def get_offloading_mask(ue_params, tasks):
    n_ues = len(ue_params)
    offloading_mask = np.zeros(n_ues, dtype=bool)

    # Compute local execution times and energies
    D_loc_list = np.array([
        task['cpu_cycles'] / ue['f_loc'] for task, ue in zip(tasks, ue_params)
    ])
    E_loc_list = np.array([
        task['cpu_cycles'] * ue['delta_loc']
        for task, ue in zip(tasks, ue_params)
    ])

    for i, (task, ue) in enumerate(zip(tasks, ue_params)):
        if D_loc_list[i] > task[
                'max_delay']:  #if the execution time for local processing is greater than the maximum delay, then we will automatically offload the task
            offloading_mask[i] = True

    return offloading_mask


class EnergyComputation:

    def __init__(self, channel_params, MEC_params):

        self.W = channel_params['bandwidth']
        self.N0 = channel_params['noise_power']
        self.f_MEC = MEC_params['f_MEC']
        self.P_MEC = MEC_params['P_MEC']

    def get_total_energy(self, offload_vector, theta_vector, beta_vector,
                         ue_params, tasks):

        total_energy_consumption = 0.0
        for i, (offload_decision, th_i, beta_i) in enumerate(
                zip(offload_vector, theta_vector, beta_vector)):
            ue = ue_params[i]
            task = tasks[i]

            D_loc = task['cpu_cycles'] / ue['f_loc']  #local execution time
            E_loc = task['cpu_cycles'] * ue[
                'delta_loc']  #local execution energy

            p_i = ue['p_trans']  # Transmission power
            g_i = ue['channel_gain']
            p_idle = ue['p_idle']

            SNR = (p_i * g_i) / (th_i * self.W * self.N0)
            R_i = th_i * self.W * np.log2(1 + SNR)

            D_trans = task['data_size'] / R_i
            D_exec = task['cpu_cycles'] / (beta_i * self.f_MEC)
            E_trans = p_i * D_trans
            E_MEC = self.P_MEC * beta_i * D_exec
            E_idle = p_idle * D_exec

            D_offload = D_trans + D_exec
            E_offload = E_trans + E_MEC + E_idle

            if offload_decision == True:
                total_energy_consumption += E_offload
            else:
                total_energy_consumption += E_loc

        return total_energy_consumption


def run_DQN():
    gamma = 0.99
    lr = 0.00005
    batch_size = 32
    memory_size = 500
    target_update_freq = 50

    epsilon = epsilon_start = 1.0
    epsilon_end = 0.1
    epsilon_decay = 0.995

    n_ues = 3

    energy_computation = EnergyComputation(channel_params(), MEC_params())

    action_space = ActionSpace(n_ues)  #discretized action space
    n_actions = action_space.n_actions

    state_dim = 3

    dqn_model = CNN_DQN(state_dim, n_actions).to(device)
    target_model = CNN_DQN(state_dim, n_actions).to(device)
    target_model.load_state_dict(dqn_model.state_dict())  # copy weights
    optimizer = optim.Adam(dqn_model.parameters(), lr=lr)

    n_episodes = 150
    n_timesteps = 500

    all_total_energies = []
    all_rewards = []
    all_losses = []

    for episode in range(n_episodes):
        mean_total_energy = []
        mean_reward = []
        mean_loss = []

        memory = ReplayBuffer(memory_size)

        #state space: total energy consumed, available computation resources, availabile spectrum resources
        initial_total_energy_consumption = 0.0
        initial_available_computation = 1.0
        initial_available_spectrum = 1.0

        initial_state = torch.tensor([
            initial_total_energy_consumption, initial_available_computation,
            initial_available_spectrum
        ],
                                     dtype=torch.float32,
                                     device=device)

        current_state = initial_state.clone().to(device)
        previous_task_info = torch.zeros_like(initial_state,
                                              dtype=torch.float32).to(device)

        for t in range(n_timesteps):

            ue_params = ue_sample(n_ues)
            tasks = task_sample(n_ues)

            total_energy_from_state = current_state[0].item(
            )  # lingering energy from previous tasks

            if random.random() < epsilon:
                action = random.randint(0, n_actions - 1)
            else:
                with torch.no_grad():
                    q_values = dqn_model(current_state)
                    action = torch.argmax(q_values).item()

            offload_vector, beta_vector, theta_vector = action_space.decode(
                action)
            new_task_energy = energy_computation.get_total_energy(
                offload_vector, theta_vector, beta_vector, ue_params, tasks)

            total_energy_consumption = total_energy_from_state + new_task_energy

            mean_total_energy.append(total_energy_consumption)

            if sum(beta_vector) > 1 or sum(theta_vector) > 1:
                reward = -2 * (total_energy_consumption)
            else:
                reward = -total_energy_consumption

            mean_reward.append(reward)

            beta_sum = min(sum(beta_vector), 1.0)
            theta_sum = min(sum(theta_vector), 1.0)

            current_task_info = torch.tensor(
                [new_task_energy, 1 - beta_sum, 1 - theta_sum],
                dtype=torch.float,
                device=device)

            next_state = current_state + current_task_info - previous_task_info

            memory.push(current_state.clone(), action, reward,
                        next_state.clone())

            previous_task_info = current_task_info.clone()
            current_state = next_state.clone()

            if len(memory) >= batch_size:
                states, actions, rewards, next_states = memory.sample(
                    batch_size)
                q_pred = dqn_model(states).gather(
                    1, actions.unsqueeze(1)).squeeze(1)
                with torch.no_grad():
                    next_actions = dqn_model(next_states).argmax(dim=1)
                    q_next = target_model(next_states).gather(
                        1, next_actions.unsqueeze(1)).squeeze(1)
                q_target = rewards + gamma * q_next

                #loss = torch.nn.functional.smooth_l1_loss(q_pred, q_target)
                loss = nn.MSELoss()(q_pred, q_target)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(dqn_model.parameters(), 1.0)
                optimizer.step()

                mean_loss.append(loss.detach().item())

            if t % target_update_freq == 0:
                target_model.load_state_dict(dqn_model.state_dict())

        epsilon = max(epsilon_end, epsilon * epsilon_decay)

        print(max(mean_total_energy), ' min ', min(mean_total_energy))
        episode_mean_energy = sum(mean_total_energy) / len(mean_total_energy)
        episode_mean_reward = sum(mean_reward) / len(mean_reward)
        episode_mean_loss = sum(mean_loss) / len(mean_loss)

        print(
            f"Episode {episode+1}/{n_episodes} : Mean Total Energy: {episode_mean_energy:.4f} , Mean Reward: {episode_mean_reward:.4f}, Mean Loss: {episode_mean_loss:.4f}"
        )

        all_total_energies.append(episode_mean_energy)
        all_rewards.append(episode_mean_reward)
        all_losses.append(episode_mean_loss)

    return all_total_energies, all_rewards, all_losses


if __name__ == "__main__":

    experiment_name = "CNN-Based DDQN"
    run_save_and_plot_experiment(experiment_name, n_runs=5)