in the below code, the actor network using the "gnn_model", where as the critic network is not using, is it good implementation, can i proceed my computation as it is? or should i need to include the same "gnn_model" also in the critic network?

# Instantiate the GNN model
gnn_model = CustomGNN(in_channels, out_channels)

agent = PPOAgent(gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef)

class CustomGNN(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(CustomGNN, self).__init__()
        self.gat1 = GATConv(in_channels, 8, heads=8, dropout=0.2) # dropout=0.6
        self.gat2 = GATConv(64, out_channels, heads=1, concat=False, dropout=0.2)  # 8*8 is 64 # dropout=0.6

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.2, training=self.training) # p=0.6
        x = F.elu(self.gat1(x, edge_index))
        # After dropout and activation
        #assert torch.equal(X1[:, 0], x[:, 0]), "Order mismatch after dropout/activation."
        x = F.dropout(x, p=0.2, training=self.training) # p=0.6
        x = self.gat2(x, edge_index)

        return x

class Actor(torch.nn.Module):
    def __init__(self, gnn_model):
        super(Actor, self).__init__()
        self.gnn = gnn_model
        self.stability_head = nn.Sequential(
            nn.Linear(out_channels, 64),  # Adjusted input to match GNN output features
            nn.ReLU(),
            nn.Linear(64, 1)    # Outputs a single stability value per node
        )
        self.action_head = nn.Linear(out_channels, action_space)  # Map to action space
        
        # Define bounds for action space scaling
        self.bounds_low = torch.tensor([0.18e-6, 0.18e-6, 0.18e-6, 0.18e-6, 0.18e-6, 
                                        0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 
                                        15e-6, 0.1e-12, 0.8], dtype=torch.float32)
                                        
        self.bounds_high = torch.tensor([0.181e-6, 0.181e-6, 0.181e-6, 0.181e-6, 0.181e-6, 
                                         50e-6, 50e-6, 50e-6, 50e-6, 50e-6, 
                                         30e-6, 10e-12, 1.4], dtype=torch.float32)
        
    def forward(self, state):
        node_features_tensor, _, edge_index = state
        
        # Processed features from the GNN
        processed_features = self.gnn(node_features_tensor, edge_index)        
        # Example assertion (add after GNN processing)
        #assert torch.all(node_features_tensor[:, 0] == processed_features[:, 0]), "Node order mismatch!"
        
        # Global mean pooling to aggregate node features
        graph_embedding = gmp(processed_features, batch=torch.zeros(processed_features.size(0), dtype=torch.long, device=processed_features.device))
        
        # Initialize stability predictions - assume all nodes as stable
        stability_pred = torch.ones(processed_features.size(0), device=processed_features.device)
        
        # Define the indices of nodes that actually require stability prediction
        # Given [0, 3, 5, 7, 10, 13, 15, 16] represent tunable/unstable nodes
        unstable_nodes_idx = torch.tensor([0, 3, 5, 7, 10, 13, 15, 16], dtype=torch.long, device=processed_features.device)
        
        # Fetch features for nodes requiring stability computation
        features_for_stability = processed_features[unstable_nodes_idx]
        
        # Compute stability predictions for these nodes
        stability_pred_unstable = self.stability_head(features_for_stability).squeeze(-1)
        
        # Update the overall stability predictions
        stability_pred[unstable_nodes_idx] = stability_pred_unstable
        
        # Pass through action head to get final action means
        action_means = self.action_head(graph_embedding)
       
        # Normalize and scale the GNN outputs to the action bounds for usable actions
        normalized_gnn_output = torch.tanh(action_means)
        action_means = self.bounds_low + (self.bounds_high - self.bounds_low) * ((normalized_gnn_output + 1) / 2)
                
        # Define action variances (assuming a simple fixed form here for demonstration)
        action_log_std = torch.zeros_like(action_means)
        action_std = torch.exp(action_log_std)

        return action_means, action_std, stability_pred

class Critic(torch.nn.Module):
    def __init__(self):
        super(Critic, self).__init__()
        flattened_size = 20 * 24  # Number of nodes times features per node
        self.network = torch.nn.Sequential(
            torch.nn.Linear(flattened_size, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 1)
        )

    def forward(self, state):
        # Directly use state as it is already what we referred to as state[0] in evaluate
        flattened = state.view(-1)  # Flatten to [480]
        #node_features_tensor, _, _ = state
        #flattened = node_features_tensor.view(-1)  # Flatten to [480]
        return self.network(flattened)

class PPOAgent:
    def __init__(self, gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef):
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.epsilon = epsilon
        self.policy_clip = policy_clip
        self.epochs = epochs
        self.entropy_coef = entropy_coef  # Dynamic entropy coefficient
        
        self.actor = Actor(gnn_model)
        self.critic = Critic()
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr_critic)
        self.action_space = action_space  # Assume continuous
        
        self.log_dir = 'OPAMP_GNNRL_log_file25_3'
        os.makedirs(self.log_dir, exist_ok=True)
        
    def select_action(self, state, performance_metrics):               
        
        action_means, action_stds, _ = self.actor(state) 
        #print("action_means", action_means)

        dist = torch.distributions.Normal(action_means, action_stds)
        action = dist.sample()        
        
        log_probs = dist.log_prob(action).sum(axis=-1)  # Summing over actions if action space is multi-dimensional
        
        # Using tanh to squash gnn_output to [-1, 1]
        scaled_action = torch.tanh(action)
        
        # Now scale and shift the normalized output to the action bounds
        scaled_actions = self.actor.bounds_low + (self.actor.bounds_high - self.actor.bounds_low) * ((scaled_action + 1) / 2)

        return scaled_actions.detach().numpy(), log_probs.detach(), performance_metrics

    def update_policy(self, prev_states, prev_actions, prev_log_probs, returns, advantages):
        device = next(self.actor.parameters()).device  # Get the device from the actor model
        advantages = torch.tensor(advantages, dtype=torch.float32, device=device)
        returns = torch.tensor(returns, dtype=torch.float32, device=device)
        prev_log_probs = torch.tensor(prev_log_probs, dtype=torch.float32, device=device)        
        
        # Initialize lists to accumulate losses
        actor_losses = []
        critic_losses = []
        
        for _ in range(self.epochs):
            
            log_probs_batch, state_values_batch, entropies_batch, stability_losses_batch = [], [], [], []
    
            for state, action in zip(prev_states, prev_actions):
                node_features_tensor, _, edge_index = state
                node_features_tensor = node_features_tensor.to(device)
                edge_index = edge_index.to(device) if edge_index is not None else None
                
                # Flattening the node_features_tensor for the critic, as your critic expects a flattened input.
                flattened_state = node_features_tensor.view(-1)  # Adjusted to flatten the state appropriately
            
                # Extracting the 24th feature for stability; assuming indexing starts at 0, so it’s 23rd in 0-indexing.
                stability_target = node_features_tensor[:, 23].unsqueeze(1)  # For the stability target
    
                action_means, action_stds, stability_pred = self.actor(state)

                # Convert stability_pred to match the binary targets if not already
                stability_pred = torch.sigmoid(stability_pred).unsqueeze(1)
                
                # Getting the state value prediction from the critic using the flattened state
                state_value = self.critic(flattened_state).unsqueeze(0)  # Direct application with adjusted state