I am using PPO using Custom graph neural network as my reinforcement learning agent,
-In my implementation,I am using PPO using Custom graph neural network as my reinforcement learning agent, in the analog circuit synthesis optimization process.
-In my graph(undirected bipartite graph with circuit netlist 'components' as nodes (11 nodes) on one side and the circuit netlist 'nets' as nodes (9 nodes) on another side total 20 nodes, and the edges were connected between the components nodes and nets nodes) construction process, Fixed circuit topology (netlist) is converted in to graph structure. During optimization process, in this fixed topology(netlist/graph), the node features of circuit netlist component nodes(circuit netlist components values/dimensions) are alone dynamic in nature, and all the others are static in nature, the others are the topology of all nodes and its edges, the nodes and node features of circuit netlist 'nets', and edge features are static in nature.
-each nodes has 23 node features, among the total nodes and its node features, i need only 8 node features, each of the node features from its respective required nodes alone to my environment as input [(0, 18), (13, 18), (8, 18), (16, 18), (10, 18), (17, 20), (18, 21), (19, 22)]. based on this requirement criteria the CustomGNN class is designed.
-the circuit environment gets input action space for the variables(8) 'W1', 'W3', 'W5', 'W6', 'W7','Ib', 'Cc', 'Vc', and generate four output performance metrics 'gain', 'bw', 'ugf', 'pm' for each action space input, and from the environment we get state data(node feature tensor and edge index) and reward values(whether the output performance metrics reach the target metrics or not, based on the difference we calculate the cumulated reward values).

-From the below code i need to know why the stability head is added in the actor function, whether it is correct, or do i need it for the actor model, because i donot have any stability related node features given as input to the network. please review the code completely and provide the necessary modification which ever required in the code.
        self.stability_head = nn.Sequential(
            nn.Linear(23, 64),  # 23 features from GNN’s processed features
            nn.ReLU(),
            nn.Linear(64, 1)    # Outputs a single stability value per node
        )
# Initialize Hyperparameters
in_channels = 23  # Number of input features
out_channels = 23  # Number of output actions or features
state_dim = 23  # Dimension of state representation, to match critic input
action_space = 8  # Environment’s action space
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 
        # hidden_features = 8
        # self.gat1 = GATConv(input_features, hidden_features, heads=heads, concat=True)
        # self.gat2 = GATConv(hidden_features * heads, output_features, heads=1, concat=False)
        # hidden_features * heads = 8 * 8 = 64
        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)
        # Selective feature synchronization based on given mappings
        # Converting given (row, column) pairs to 0-based indexing
        sync_mappings = {            
            (0, 18): (5, 18),
            (0, 19): (5, 19),
            (13, 18): (15, 18),
            (13, 19): (15, 19),
            (8, 18): (12, 18),
            (8, 19): (12, 19)            
        }
        
        original_features = x.clone()
        for target, source in sync_mappings.items():
            original_features[target[0], target[1]] = original_features[source[0], source[1]]
        
        # Only retain the transformative updates for the specified 19 features
        dynamic_updates_mask = torch.zeros_like(x, dtype=bool)
        dynamic_indices = [
            (0, 18), (5, 18), (8, 18), (10, 18), (12, 18), (13, 18), (15, 18), (16, 18),
            (17, 20), (18, 21), (19, 22)]       
        
        for row, col in dynamic_indices:
            dynamic_updates_mask[row, col] = True
        
        x = original_features * (~dynamic_updates_mask) + x * dynamic_updates_mask

        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(23, 64),  # 23 features from GNN’s processed features
            nn.ReLU(),
            nn.Linear(64, 1)    # Outputs a single stability value per node
        )
        
        # Define bounds for action space scaling
        self.bounds_low = torch.tensor([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([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!"
                      
        # For action_indices handling, adapt based on the dynamic nature of processed_features
        action_indices = [(0, 18), (13, 18), (8, 18), (16, 18), (10, 18), (17, 20), (18, 21), (19, 22)]
                          
        # Convert action_indices to tensor format for easier handling
        action_indices_tensor = torch.tensor(action_indices, dtype=torch.long).t()
        selected_features = processed_features[action_indices_tensor[0], action_indices_tensor[1]]
        
        # Normalize and scale the GNN outputs to the action bounds for usable actions
        normalized_gnn_output = torch.tanh(selected_features)
        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
  
class Critic(torch.nn.Module):
    def __init__(self, gnn_model):
        super(Critic, self).__init__()
        self.gnn_model = gnn_model
        self.fc = torch.nn.Sequential(
            #torch.nn.Linear(24 * 20, 128),
            torch.nn.Linear(23 * 20, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 1)
        )

    def forward(self, node_features_tensor, edge_index):
        x = self.gnn_model(node_features_tensor, edge_index)
        x = x.view(-1)  # Flatten
        return self.fc(x)

class PPOAgent:
    def __init__(self, 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
        
        gnn_model_actor = CustomGNN(in_channels=state_dim, out_channels=23)
        gnn_model_critic = CustomGNN(in_channels=state_dim, out_channels=23)
        
        self.actor = Actor(gnn_model_actor)
        self.critic = Critic(gnn_model_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_LTS1_log_file'
        os.makedirs(self.log_dir, exist_ok=True)
        
    def save_checkpoint(self, episode):
        torch.save(self.actor.state_dict(), os.path.join(self.log_dir, f'actor_{episode}.pt'))
        torch.save(self.critic.state_dict(), os.path.join(self.log_dir, f'critic_{episode}.pt'))

    def load_checkpoint(self, episode):
        self.actor.load_state_dict(torch.load(os.path.join(self.log_dir, f'actor_{episode}.pt')))
        self.critic.load_state_dict(torch.load(os.path.join(self.log_dir, f'critic_{episode}.pt')))

    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)
        
        # Converting to numpy array for easier manipulation
        scaled_actions_np = scaled_actions.detach().numpy()
        
        # Rounding operations as per requirements
        rounded_actions = np.zeros_like(scaled_actions_np)        
        
        # Rounding to two decimal places after scaling by 1e-6 for L1, L3, L5, L6, L7, W1, W3, W5, W6, W7
        rounded_actions[:5] = np.round(scaled_actions_np[:5] * 1e6, 2) / 1e6
        
        # Rounding to the nearest whole number after scaling by 1e-6 for Ib
        rounded_actions[5] = np.round(scaled_actions_np[5] * 1e6) / 1e6
        
        # Rounding to one decimal place for Cc. Scale by 1e12, round, and scale back.
        rounded_actions[6] = np.round(scaled_actions_np[6] * 1e12, 1) / 1e12

        # Rounding to one decimal place for Vb directly
        rounded_actions[7] = np.round(scaled_actions_np[7], 1)  # for Vb
        print("rounded_actions", rounded_actions)
        return rounded_actions, log_probs.detach(), performance_metrics

    def compute_gae(self, next_value, rewards, masks, values):
        # Generalized Advantage Estimation (GAE)
        returns = []
        gae = 0
        for step in reversed(range(len(rewards))):
            delta = rewards[step] + self.gamma * next_value * masks[step] - values[step]
            gae = delta + self.gamma * self.gae_lambda * masks[step] * gae
            returns.insert(0, gae + values[step])
            next_value = values[step]
        return returns

    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 = [], [], []
    
            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                
      
                action_means, action_stds = self.actor(state)                
               
                # Getting the state value prediction from the critic using the flattened state
                state_value = self.critic(node_features_tensor, edge_index).unsqueeze(0)  # Direct application with adjusted state
    
                dist = torch.distributions.Normal(action_means, action_stds)
                action_tensor = torch.tensor(action, dtype=torch.float32, device=action_means.device)
                log_prob = dist.log_prob(action_tensor).sum(axis=-1).unsqueeze(0)
                log_probs_batch.append(log_prob)
 
                entropy = dist.entropy().sum(axis=-1).unsqueeze(0)
                entropies_batch.append(entropy)
    
                state_values_batch.append(state_value)  # No additional unsqueeze(0) needed here after correction
            
            # Convert to tensors
            log_probs = torch.cat(log_probs_batch).to(prev_log_probs.device)
            state_values = torch.cat(state_values_batch).squeeze(-1).to(returns.device)  # Ensuring the correct final shape for state_values
            entropy = torch.cat(entropies_batch)
    
            # Surrogate function
            ratios = torch.exp(log_probs - prev_log_probs)            
            #advantages_batch = returns - state_values.detach()
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1.0 - self.policy_clip, 1.0 + self.policy_clip) * advantages
            
            actor_loss = - torch.min(surr1, surr2).mean() - self.entropy_coef * entropy.mean()
            critic_loss = F.mse_loss(state_values, returns)            
            
            # Update actor
            self.optimizer_actor.zero_grad()
            #actor_loss.backward()
            actor_loss.backward(retain_graph=True)
            torch.nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=1.0)
            self.optimizer_actor.step()
            
            # Update critic
            self.optimizer_critic.zero_grad()
            critic_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=1.0)
            self.optimizer_critic.step()
            
            # Append the item loss to the lists
            actor_losses.append(actor_loss.item())  # Use item() to extract scalar value
            critic_losses.append(critic_loss.item())  # Use item() to extract scalar value
            
        # Calculate average loss over all epochs
        average_loss_actor = sum(actor_losses) / len(actor_losses)
        average_loss_critic = sum(critic_losses) / len(critic_losses)
    
        return average_loss_actor, average_loss_critic