How to implement the below suggestions, provide me the updated code in the relevane methods to satisfies the below suggestions,
- Ensuring Saturation State: While the code includes a mechanism to predict and influence the stability state, the direct connection between this prediction and ensuring the MOS transistors remain in saturation is not explicitly made in the code. This linkage depends on how well the stability_head is trained and how precisely the stability condition of a MOS transistor in saturation can be represented by the 24th node feature.
- Validation Logic: There’s no explicit validation logic in the given code snippets to check whether the updated node features indeed satisfy the stability condition before taking an action. The stability predictions are generated and influence the loss, but there’s no “hard check” before an action is considered valid.

### Recommendations:
- Explicit Checks or Constraints: Consider implementing additional layers of validation or constraints that ensure the actions proposed indeed lead to a stable state, especially if the stability condition relates to complex physical properties like MOS transistor operation modes.

class Actor(torch.nn.Module):
    def __init__(self, gnn_model):
        super(Actor, self).__init__()
        self.gnn = gnn_model
        # Bounds are converted to tensors for ease of calculation
        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.2e-6, 0.2e-6, 0.2e-6, 0.2e-6, 0.2e-6, 
                                         50e-6, 50e-6, 50e-6, 50e-6, 50e-6, 
                                         30e-6, 10e-12, 1.4], dtype=torch.float32)
        
        # Assuming num_nodes is the number of nodes and each node’s stability is predicted separately.
        self.stability_head = torch.nn.Sequential(
            torch.nn.Linear(24, 64),  # gnn_model.out_channels depends on your GNN’s output feature size
            torch.nn.ReLU(),
            torch.nn.Linear(64, 1)  # Predicting a single stability value for each node
        )
        

    def forward(self, state):
        node_features_tensor, _, edge_index = state
        processed_features = self.gnn(node_features_tensor, edge_index)
        
        # Specific (row, column) indices for action values, converted to 0-based indexing
        action_indices = [
            (10, 19), (16, 19), (5, 19), (3, 19), (0, 19), (10, 18),
            (16, 18), (5, 18), (3, 18), (0, 18), (17, 20), (18, 21), (19, 22)
        ]
        
        # For gathering specific indices, we first convert action_indices to a tensor format that can be used with gather or indexing
        action_indices_tensor = torch.tensor(action_indices, dtype=torch.long).t()
        selected_features = processed_features[action_indices_tensor[0], action_indices_tensor[1]]
        
        # Using tanh to squash gnn_output to [-1, 1]
        normalized_gnn_output = torch.tanh(selected_features)
        
        # Now scale and shift the normalized output to the action bounds
        action_means = self.bounds_low + (self.bounds_high - self.bounds_low) * ((normalized_gnn_output + 1) / 2)

        # Here we should define action variances as well, which are learnable and specific to your model structure
        action_log_std = torch.zeros_like(action_means)
        action_std = torch.exp(action_log_std)
        
        # Predict stability for each node, assuming processed_features has size [num_nodes, out_channels]
        stability_pred = self.stability_head(processed_features).squeeze(-1)  # Ensure compatibility: [num_nodes] expected
        #print("stability_pred", stability_pred)
        return action_means, action_std, stability_pred

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

    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):
        advantages = torch.tensor(advantages, dtype=torch.float32)
        returns = torch.tensor(returns, dtype=torch.float32)
        prev_log_probs = torch.tensor(prev_log_probs, dtype=torch.float32)        
        
        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
                
                # 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)
                stability_pred = stability_pred.unsqueeze(1)  # Ensure matching shape with stability_target
                
                # 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
    
                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)
                
                stability_loss = F.binary_cross_entropy_with_logits(stability_pred, stability_target)
                stability_losses_batch.append(stability_loss.unsqueeze(0))
    
                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
            
            log_probs = torch.cat(log_probs_batch)
            state_values = torch.cat(state_values_batch).squeeze(-1)  # Ensuring the correct final shape for state_values
            stability_losses = torch.cat(stability_losses_batch)
            entropy = torch.cat(entropies_batch)
    
            
            ratios = torch.exp(log_probs - prev_log_probs.detach())            
            advantages_batch = returns - state_values.detach()
            surr1 = ratios * advantages_batch
            surr2 = torch.clamp(ratios, 1.0 - self.policy_clip, 1.0 + self.policy_clip) * advantages_batch
            
            actor_loss = - torch.min(surr1, surr2).mean() - self.entropy_coef * entropy.mean() + stability_losses.mean()
            critic_loss = F.mse_loss(state_values, returns)            
           
            self.optimizer_actor.zero_grad()
            actor_loss.backward()
            self.optimizer_actor.step()
            
            self.optimizer_critic.zero_grad()
            critic_loss.backward()
            self.optimizer_critic.step()