I have an idea, this is my current model:"from . import RLAgent
from common.registry import Registry
import numpy as np
import os
import random
from collections import OrderedDict, deque
import gym

from generator.lane_vehicle import LaneVehicleGenerator
from generator.intersection_phase import IntersectionPhaseGenerator
import torch
from torch import nn
import torch.nn.functional as F
import torch_scatter
import torch.optim as optim
from torch.nn.utils import clip_grad_norm_

from torch_geometric.nn import MessagePassing
from torch_geometric.data import Data, Batch
from torch_geometric.utils import add_self_loops

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
torch.autograd.set_detect_anomaly(True)

@Registry.register_model('TestLight')
class TestLightAgent(RLAgent):
    def __init__(self, world, rank):
        super().__init__(world, world.intersection_ids[rank])
        self.rank = rank  # Initialize the rank attribute
        self.buffer_size = Registry.mapping['trainer_mapping']['setting'].param['buffer_size']
        self.replay_buffer = deque(maxlen=self.buffer_size)

        self.graph = Registry.mapping['world_mapping']['graph_setting'].graph
        self.world = world
        self.sub_agents = len(self.world.intersections)
        self.edge_idx = torch.tensor(self.graph['sparse_adj'].T, dtype=torch.long).to(device)
        
        self.phase = Registry.mapping['model_mapping']['setting'].param['phase']
        self.one_hot = Registry.mapping['model_mapping']['setting'].param['one_hot']
        self.model_dict = Registry.mapping['model_mapping']['setting'].param

        # Initialize intersection importance with ones for all nodes
        self.intersection_importance = {i: torch.tensor(1.0, device=device) for i in range(self.sub_agents)}

        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ['lane_count'], in_only=True, average=None)
            observation_generators.append((node_idx, tmp_generator))
        observation_generators.sort(key=lambda x: x[0])  # Corrected sorting method
        self.ob_generator = observation_generators

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"],
                                                 in_only=True, average='all', negative=True)
            rewarding_generators.append((node_idx, tmp_generator))
        rewarding_generators.sort(key=lambda x: x[0])
        self.reward_generator = rewarding_generators

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"], 
                                                 in_only=True, negative=False)
            queues.append((node_idx, tmp_generator))
        queues.sort(key=lambda x: x[0])
        self.queue = queues

        delays = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_delay"], 
                                                 in_only=True, average="all", negative=False)
            delays.append((node_idx, tmp_generator))
        delays.sort(key=lambda x: x[0])
        self.delay = delays

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, inter, ['phase'],
                                                       targets=['cur_phase'], negative=False)
            phasing_generators.append((node_idx, tmp_generator))
        phasing_generators.sort(key=lambda x: x[0])
        self.phase_generator = phasing_generators

        self.action_space = gym.spaces.Discrete(len(self.world.intersections[0].phases))
        min_ob_length = max([ob[1].ob_length for ob in self.ob_generator])
        self.ob_length = min_ob_length


        self.learning_rate = Registry.mapping['model_mapping']['setting'].param['learning_rate']

        self.target_update_freq = Registry.mapping['trainer_mapping']['setting'].param['target_update_freq']

        self.gamma = Registry.mapping['model_mapping']['setting'].param['gamma']
        self.grad_clip = Registry.mapping['model_mapping']['setting'].param['grad_clip']
        self.epsilon = Registry.mapping['model_mapping']['setting'].param['epsilon']
        self.epsilon_decay = Registry.mapping['model_mapping']['setting'].param['epsilon_decay']
        self.epsilon_min = Registry.mapping['model_mapping']['setting'].param['epsilon_min']

        self.vehicle_max = Registry.mapping['model_mapping']['setting'].param['vehicle_max']
        self.batch_size = Registry.mapping['model_mapping']['setting'].param['batch_size']

        self.episode_count = 0
        self.update_interval = Registry.mapping['trainer_mapping']['setting'].param['update_interval']  # Initial interval for updating importance
        self.decay_factor = Registry.mapping['trainer_mapping']['setting'].param['decay_factor']  # Factor to increase the interval after each update
        self.softmax_temperature = Registry.mapping['trainer_mapping']['setting'].param['softmax_temperature']  # Temperature for softmax
        self.softmax_feather = Registry.mapping['trainer_mapping']['setting'].param['softmax_feather']  # Feather for softmax
    
        self.model = self._build_model().to(device)
        self.target_model = self._build_model().to(device)
        self.optimizer = optim.AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=1e-5)
        self.criterion = nn.MSELoss()

    def reset(self):
        # Similar to __init__ but for resetting generators
        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ['lane_count'], in_only=True, average=None)
            observation_generators.append((node_idx, tmp_generator))
        observation_generators.sort(key=lambda x: x[0])
        self.ob_generator = observation_generators

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"],
                                                 in_only=True, average='all', negative=True)
            rewarding_generators.append((node_idx, tmp_generator))
        rewarding_generators.sort(key=lambda x: x[0])
        self.reward_generator = rewarding_generators

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, inter, ['phase'],
                                                       targets=['cur_phase'], negative=False)
            phasing_generators.append((node_idx, tmp_generator))
        phasing_generators.sort(key=lambda x: x[0])
        self.phase_generator = phasing_generators

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"], 
                                                 in_only=True, negative=False)
            queues.append((node_idx, tmp_generator))
        queues.sort(key=lambda x: x[0])
        self.queue = queues

        delays = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_delay"], 
                                                 in_only=True, average="all", negative=False)
            delays.append((node_idx, tmp_generator))
        delays.sort(key=lambda x: x[0])
        self.delay = delays

    def get_ob(self):
        x_obs = []  # sub_agents * lane_nums,
        for i in range(len(self.ob_generator)):
            ob = self.ob_generator[i][1].generate() / self.vehicle_max
            ob = np.pad(ob, (0, self.ob_length - ob.shape[-1]))
            x_obs.append(ob)
        x_obs = np.array(x_obs, dtype=np.float32)
        return x_obs

    def get_reward(self):
        rewards = []
        for i in range(len(self.reward_generator)):
            rewards.append(self.reward_generator[i][1].generate())
        rewards = np.squeeze(np.array(rewards, dtype=np.float32)) * 12
        return rewards

    def get_phase(self):
        phase = []
        for i in range(len(self.phase_generator)):
            phase.append((self.phase_generator[i][1].generate()))
        phase = (np.concatenate(phase)).astype(np.int8)
        return phase

    def get_queue(self):
        queue = []
        for item in self.queue:
            item = item[1].generate()
            item = np.pad(item, (0, self.ob_length - item.shape[-1]))
            queue.append(item)
        tmp_queue = np.squeeze(np.array(queue, dtype=np.float32))
        queue = np.sum(tmp_queue, axis=1 if len(tmp_queue.shape) == 2 else 0)
        return queue

    def get_delay(self):
        delay = []
        for i in range(len(self.delay)):
            delay.append((self.delay[i][1].generate()))
        delay = np.squeeze(np.array(delay, dtype=np.float32))
        return delay

    def get_action(self, ob, phase, test=False):
        if not test and np.random.rand() <= self.epsilon:
            return self.sample()

        observation = torch.tensor(ob, dtype=torch.float32).to(device)
        edge_index = self.edge_idx.to(device)

        # Corrected: Generate mask based on the number of nodes in the observation
        tgt_mask = self.generate_square_subsequent_mask(observation.size(0)).to(device)
        memory_mask = None  # Define as needed

        q_values, attention_matrices = self.model(
            x=observation, 
            edge_index=edge_index, 
            train=False, 
            tgt_mask=tgt_mask, 
            memory_mask=memory_mask
        )
        actions = torch.argmax(q_values, dim=1)

        self.update_intersection_importance(attention_matrices)

        return actions.cpu().numpy()

    def sample(self):
        return np.random.randint(0, self.action_space.n, self.sub_agents)

    def generate_square_subsequent_mask(self, sz):
        mask = torch.triu(torch.ones(sz, sz) * float('-inf'), diagonal=1)
        return mask

    def update_intersection_importance(self, attention_matrices):
        importance = {}
        last_attention_matrix = attention_matrices[-1]  # Use only the last attention matrix

        # Calculate raw importance scores
        for i in range(self.sub_agents):
            if i < last_attention_matrix.size(0):
                importance[i] = last_attention_matrix[i].mean().detach()
            else:
                # If we don't have attention for this node, use the previous importance or a default value
                importance[i] = self.intersection_importance.get(i, torch.tensor(0.5, device=device))

        # Convert importance dict to tensor
        importance_tensor = torch.tensor([importance[i] for i in range(self.sub_agents)], device=device)

        # Apply softmax with a lower temperature to make the distribution sharper
        temperature = self.softmax_temperature  # Lower value for sharper distribution
        normalized_importance_tensor = F.softmax(importance_tensor / temperature, dim=0)

        # Convert back to dictionary
        normalized_importance = {i: normalized_importance_tensor[i].item() for i in range(self.sub_agents)}

        self.intersection_importance = normalized_importance
        self.model.intersection_importance = normalized_importance
        self.target_model.intersection_importance = normalized_importance
        #PRINT DEBUG STATEMENT
        #print(f"Updated Normalized Intersection Importance: {self.intersection_importance}")
        #print(f"Number of nodes: {self.sub_agents}")
        #print(f"Size of importance dict: {len(self.intersection_importance)}")
        #print(f"Sum of importance scores: {sum(self.intersection_importance.values())}")

    def _build_model(self):
        model = TestLightNet(
            self.ob_length, 
            self.action_space.n, 
            self.intersection_importance, 
            softmax_temperature=self.softmax_temperature,
            softmax_feather=self.softmax_feather,
            update_interval=self.update_interval,
            decay_factor=self.decay_factor,
            **self.model_dict
        )
        return model

    def remember(self, last_obs, last_phase, actions, actions_prob, rewards, obs, cur_phase, done, key):
        self.replay_buffer.append((key, (last_obs, last_phase, actions, rewards, obs, cur_phase)))

    def _batchwise(self, samples):
        batch_list = []
        batch_list_p = []
        actions = []
        rewards = []
        for item in samples:
            dp = item[1]
            state = torch.tensor(dp[0], dtype=torch.float32)
            batch_list.append(Data(x=state, edge_index=self.edge_idx))

            state_p = torch.tensor(dp[4], dtype=torch.float32)
            batch_list_p.append(Data(x=state_p, edge_index=self.edge_idx))
            rewards.append(dp[3])
            actions.append(dp[2])
        batch_t = Batch.from_data_list(batch_list)
        batch_tp = Batch.from_data_list(batch_list_p)
        rewards = torch.tensor(np.array(rewards), dtype=torch.float32)
        actions = torch.tensor(np.array(actions), dtype=torch.long)
        if self.sub_agents > 1:
            rewards = rewards.view(rewards.shape[0] * rewards.shape[1])
            actions = actions.view(actions.shape[0] * actions.shape[1])
        return batch_t, batch_tp, rewards, actions

    def train(self):
        if len(self.replay_buffer) < self.batch_size:
            return 0

        samples = random.sample(self.replay_buffer, self.batch_size)
        b_t, b_tp, rewards, actions = self._batchwise(samples)

        # Ensure all tensors are on the correct device
        b_t = b_t.to(device)
        b_tp = b_tp.to(device)
        rewards = rewards.to(device)
        actions = actions.to(device)

        # Corrected: Generate mask based on the total number of nodes in the batch
        tgt_mask = self.generate_square_subsequent_mask(b_t.num_nodes).to(device)
        memory_mask = None  # Define as needed

        current_q_values, current_attention = self.model(
            x=b_t.x, 
            edge_index=b_t.edge_index, 
            train=True, 
            tgt_mask=tgt_mask, 
            memory_mask=memory_mask)
            
        target_q_values, _ = self.target_model(x=b_tp.x, edge_index=b_tp.edge_index, train=False, tgt_mask=None, memory_mask=None)

        target_q_values = rewards + self.gamma * torch.max(target_q_values, dim=1)[0]

        current_q_values = current_q_values.gather(1, actions.unsqueeze(1)).squeeze(1)

        loss = self.criterion(current_q_values, target_q_values)
        self.optimizer.zero_grad()
        loss.backward()
        clip_grad_norm_(self.model.parameters(), self.grad_clip)
        self.optimizer.step()

        # Update episode count
        self.episode_count += 1

        # Update intersection importance based on the episode count
        if self.episode_count >= self.update_interval:
            self.update_intersection_importance(current_attention)
            self.update_interval = int(self.update_interval * self.decay_factor)  # Increase the interval
            self.episode_count = 0  # Reset episode count

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

        return loss.item()

    def update_target_network(self):
        weights = self.model.state_dict()
        self.target_model.load_state_dict(weights)

    def load_model(self, e):
        model_name = os.path.join(Registry.mapping['logger_mapping']['path'].path,
                                  'model', f'{e}_{self.rank}.pt')        
        self.model = self._build_model()
        self.model.load_state_dict(torch.load(model_name))
        self.target_model = self._build_model()
        self.target_model.load_state_dict(torch.load(model_name))

    def save_model(self, e):
        path = os.path.join(Registry.mapping['logger_mapping']['path'].path, 'model')
        if not os.path.exists(path):
            os.makedirs(path)
        model_name = os.path.join(path, f'{e}_{self.rank}.pt')
        torch.save(self.target_model.state_dict(), model_name)


class TestLightNet(nn.Module):
    def __init__(self, input_dim, output_dim, intersection_importance, **kwargs):
        super(TestLightNet, self).__init__()
        self.model_dict = kwargs
        self.action_space = gym.spaces.Discrete(output_dim)
        self.features = input_dim
        self.intersection_importance = intersection_importance
        self.softmax_temperature = kwargs.get('softmax_temperature', 1.0)
        self.softmax_feather = kwargs.get('softmax_feather', 1.0)
        self.update_interval = kwargs.get('update_interval', 100)
        self.decay_factor = kwargs.get('decay_factor', 1.05)
        
        # Embedding MLP
        self.embedding_MLP = Embedding_MLP(self.features, layers=self.model_dict.get('NODE_EMB_DIM'))

        self.embedding_output_dim = self.model_dict.get('NODE_EMB_DIM')[-1]

        # Lower Transformer (Encoder)
        self.num_lower_transformer_layers = self.model_dict.get('num_lower_transformer_layers', 1)
        self.lower_transformer_nhead = self.model_dict.get('lower_transformer_nhead', 4)
        self.lower_transformer_dim_feedforward = self.model_dict.get('lower_transformer_dim_feedforward', 1024)
        self.lower_transformer_dropout = self.model_dict.get('lower_transformer_dropout', 0.1)

        self.lower_transformer = TransformerBlock(
            d_model=self.embedding_output_dim,
            nhead=self.lower_transformer_nhead,
            num_layers=self.num_lower_transformer_layers,
            dim_feedforward=self.lower_transformer_dim_feedforward,
            dropout=self.lower_transformer_dropout,
            is_encoder=True  # Initialize as encoder
        )

        # GAT Layers
        self.num_gat_layers = self.model_dict.get('N_LAYERS', 5)
        self.num_heads = self.model_dict.get('NUM_HEADS', [1])[0]

        self.head_dim = self.embedding_output_dim // self.num_heads

        self.gat_layers = nn.ModuleList([
            MultiHeadAttModel(d=self.embedding_output_dim, dv=self.head_dim, d_out=self.embedding_output_dim, nv=self.num_heads, suffix=f'layer_{i}')
            for i in range(self.num_gat_layers)
        ])

        self.skip_connections = nn.ModuleList([
            nn.Linear(self.embedding_output_dim, self.embedding_output_dim)
            for _ in range(self.num_gat_layers)
        ])

        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(self.embedding_output_dim)
            for _ in range(self.num_gat_layers)
        ])

        # Upper Transformer (Decoder)
        self.num_upper_transformer_layers = self.model_dict.get('num_upper_transformer_layers', 1)
        self.upper_transformer_nhead = self.model_dict.get('upper_transformer_nhead', 4)
        self.upper_transformer_dim_feedforward = self.model_dict.get('upper_transformer_dim_feedforward', 1024)
        self.upper_transformer_dropout = self.model_dict.get('upper_transformer_dropout', 0.1)

        self.upper_transformer = TransformerBlock(
            d_model=self.embedding_output_dim,
            nhead=self.upper_transformer_nhead,
            num_layers=self.num_upper_transformer_layers,
            dim_feedforward=self.upper_transformer_dim_feedforward,
            dropout=self.upper_transformer_dropout,
            is_encoder=False  # Initialize as decoder
        )

        # Output layer
        self.output_layer = nn.Sequential(
            nn.Linear(self.embedding_output_dim, 128),  # First hidden layer
            nn.LeakyReLU(inplace=False),                      # Activation function
            nn.Linear(128, self.action_space.n)        # Output layer
        )
        self.relu = nn.ReLU(inplace=False)
        self.batch_norm = nn.BatchNorm1d(self.embedding_output_dim)
        print(f"Hyperparameters: num_layers={self.num_lower_transformer_layers},")
        print(f"dim_feedforward={self.lower_transformer_dim_feedforward},")
        print(f"dropout={self.lower_transformer_dropout},")
        print(f"num_gat_layers={self.num_gat_layers},")
        print(f"num_heads={self.num_heads},")
        print(f"upper_transformer_nhead={self.upper_transformer_nhead},")
        print(f"upper_transformer_dim_feedforward={self.upper_transformer_dim_feedforward},")
        print(f"upper_transformer_dropout={self.upper_transformer_dropout},")
        print(f"softmax_temperature={self.softmax_temperature},")
        print(f"softmax_feather={self.softmax_feather},")
        print(f"update_interval={self.model_dict.get('update_interval', 10)},")
        print(f"decay_factor={self.model_dict.get('decay_factor', 1.005)}")

    def forward(self, x, edge_index, train=True, tgt_mask=None, memory_mask=None):
        #(f"Input x shape: {x.shape}")
        #print(f"Edge index shape: {edge_index.shape}")

        h = self.embedding_MLP(x, train)
        h = self.batch_norm(h)
        h = h.unsqueeze(0)  # Shape: (1, num_nodes, embedding_dim)

        # Prepare memory for Decoder (use lower transformer as memory)
        memory = self.lower_transformer(h)  # Encoder part
        memory = memory.squeeze(0)  # Shape: (num_nodes, embedding_dim)
        h = h.squeeze(0)  # Back to (num_nodes, embedding_dim)

        # GAT layers
        attention_matrices = []
        importance_scores = torch.tensor([self.intersection_importance.get(i, 0.5) for i in range(h.size(0))], dtype=torch.float32, device=h.device)

        for layer_idx in range(self.num_gat_layers):
            h_new, att = self.gat_layers[layer_idx](h, edge_index, train)
            h_skip = self.skip_connections[layer_idx](h)

            # Blend new and skip connections based on importance score
            blend_factor = importance_scores.unsqueeze(1) * ((layer_idx + 1) / self.num_gat_layers)
            h = blend_factor * h_new + (1 - blend_factor) * h_skip

            h = self.layer_norms[layer_idx](h)
            h = self.relu(h)

            attention_matrices.append(att)

        #h_skip = h.clone()
        h = h.unsqueeze(0)  # Shape: (1, num_nodes, embedding_dim)
        memory = memory.unsqueeze(0)  # Shape: (1, num_nodes, embedding_dim)

        # Decoder Transformer using h as target and memory from lower transformer
        h = self.upper_transformer(h, memory=memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
        h = h.squeeze(0)

        # Proceed with output_layer
        h_importance = h * (1 + importance_scores.unsqueeze(1)) / self.softmax_feather
        #concat_h = torch.cat([h_importance, h_skip], dim=1)
        q_values = self.output_layer(h_importance)

        return q_values, attention_matrices


class Embedding_MLP(nn.Module):
    def __init__(self, in_size, layers):
        super(Embedding_MLP, self).__init__()
        constructor_dict = OrderedDict()
        for l_idx, l_size in enumerate(layers):
            name = f"node_embedding_{l_idx}"
            if l_idx == 0:
                h = nn.Linear(in_size, l_size)
                constructor_dict.update({name: h})
            else:
                h = nn.Linear(layers[l_idx - 1], l_size)
                constructor_dict.update({name: h})
            name = f"n_relu_{l_idx}"
            constructor_dict.update({name: nn.ReLU()})
        self.embedding_node = nn.Sequential(constructor_dict)

    def _forward(self, x):
        x = self.embedding_node(x)
        return x

    def forward(self, x, train=True):
        if train:
            return self._forward(x)
        else:
            with torch.no_grad():
                return self._forward(x)


class MultiHeadAttModel(MessagePassing):
    def __init__(self, d, dv, d_out, nv, suffix):
        super(MultiHeadAttModel, self).__init__(aggr='add')
        self.d = d
        self.dv = dv
        self.d_out = d_out
        self.nv = nv
        self.suffix = suffix
        self.W_target = nn.Linear(d, dv * nv)
        self.W_source = nn.Linear(d, dv * nv)
        self.hidden_embedding = nn.Linear(d, dv * nv)
        self.out = nn.Linear(dv, d_out)  # nn.Linear(51, 256)
        self.att = None
        self.layer_norm = nn.LayerNorm(d_out)

    def forward(self, x, edge_index, train=True):
        if train:
            return self._forward(x, edge_index)
        else:
            with torch.no_grad():
                return self._forward(x, edge_index)

    def _forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index=edge_index)
        aggregated = self.propagate(edge_index=edge_index, x=x)
        out = self.out(aggregated)
        out = F.relu(out)
        out = self.layer_norm(out)  # Apply layer normalization
        return out, self.att

    def message(self, x_i, x_j, edge_index_i):
        h_target = F.relu(self.W_target(x_i))
        h_target = h_target.view(-1, self.nv, self.dv)  # Shape: [E, nv, dv] => [E, 5, 51]
        agent_repr = h_target.permute(1, 0, 2)  # Shape: [nv, E, dv] => [5, E, 51]

        h_source = F.relu(self.W_source(x_j))
        h_source = h_source.view(-1, self.nv, self.dv)  # Shape: [E, nv, dv] => [E, 5, 51]
        neighbor_repr = h_source.permute(1, 0, 2)  # Shape: [nv, E, dv] => [5, E, 51]

        index = edge_index_i

        e_i = torch.mul(agent_repr, neighbor_repr).sum(-1)  # Shape: [nv, E]
        max_node = torch_scatter.scatter_max(e_i, index=index)[0]  # Shape: [num_nodes,]
        max_i = max_node[:, index]  # Shape: [nv, E]

        ec_i = e_i - max_i
        ecexp_i = torch.exp(ec_i)

        norm_node = torch_scatter.scatter_sum(ecexp_i, index=index)  # Shape: [num_nodes,]
        norm_node = norm_node[:, index]
        normst_i = norm_node + 1e-12

        alpha_i = ecexp_i / normst_i  # Shape: [nv, E]
        alpha_i_expand = alpha_i.unsqueeze(-1)  # Shape: [nv, E, 1]

        hidden_neighbor = F.relu(self.hidden_embedding(x_j))  # Shape: [E, d_out]
        hidden_neighbor = hidden_neighbor.view(-1, self.nv, self.dv)  # Shape: [E, nv, dv]
        hidden_neighbor_repr = hidden_neighbor.permute(1, 0, 2)  # Shape: [nv, E, dv]

        # Element-wise multiplication and averaging
        out = (hidden_neighbor_repr * alpha_i_expand).mean(0)  # Resulting Shape: [E, dv]

        self.att = alpha_i
        return out

    def get_att(self):
        if self.att is None:
            print('invalid att')
        return self.att


class TransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, num_layers, dim_feedforward=1024, dropout=0.3, is_encoder=True):
        super(TransformerBlock, self).__init__()
        
        if is_encoder:
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                batch_first=True
            )
            self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        else:
            decoder_layer = nn.TransformerDecoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                batch_first=True
            )
            self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.is_encoder = is_encoder
    
    def forward(self, input, memory=None, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        if self.is_encoder:
            # For encoder, memory is not used
            output = self.transformer(
                src=input, 
                mask=tgt_mask,
                src_key_padding_mask=tgt_key_padding_mask
            )
            return output
        else:
            # For decoder, input is tgt, and memory comes from encoder
            output = self.transformer(
                tgt=input,
                memory=memory,
                tgt_mask=tgt_mask,
                memory_mask=memory_mask,
                tgt_key_padding_mask=tgt_key_padding_mask,
                memory_key_padding_mask=memory_key_padding_mask
            )
            return output"

new idea: Add positional embedding into the current embeddings, then split it into two flows. First flow is called GAT flow, which has all of the current GAT layer and its quirks with it. Second flow is called transformer flow, it is built off of decoder only transformer layer. For the decoder input, have the previous layer concatenate with the corresponding GAT layer of the first flow feed into it as input and repeat. Then first flow and second flow is concatenate and feed through a rotary attention block, feed through and MLP for actions (replace the current Linear layer).

Also, adjust the config file accordingly:"includes:
  - configs/tsc/base.yml

model:
  name: testlight
  train_model: True
  graphic: True
  vehicle_max: 1 #?? what is this for
  epsilon: 0.85
  epsilon_decay: 0.999
  epsilon_min: 0.01
  NEIGHBOR_NUM: 4
  NEIGHBOR_EDGE_NUM: 4
  N_LAYERS: 4
  INPUT_DIM: [128, 128]
  OUTPUT_DIM: [128, 128]
  NODE_EMB_DIM: [64, 128, 255]
  NUM_HEADS: [5]
  NODE_LAYER_DIMS_EACH_HEAD: [16, 16]
  OUTPUT_LAYERS: []
  num_lower_transformer_layers: 2
  lower_transformer_nhead: 5
  lower_transformer_dim_feedforward: 255
  lower_transformer_dropout: 0.5
  num_upper_transformer_layers: 2
  upper_transformer_nhead: 5
  upper_transformer_dim_feedforward: 255
  upper_transformer_dropout: 0.25

  #TODO: clean these up sometimes cuz all the old unused stuff is confusing
  one_hot: True
  phase: True

  dropout: 0.1

trainer:
  learning_start: 1000
  buffer_size: 10000
  batch_size: 64
  gamma: 0.99
  learning_rate: 0.001
  target_update_freq: 20
  grad_clip: 5.5
  update_interval: 20  # Initial interval for updating importance
  decay_factor: 1.01  # Factor to increase the interval after each update
  softmax_temperature: 0.4  # Temperature for softmax
  softmax_feather: 1.25  # Feather for softmax
# Additional configurations can be added as needed"