The below code is designed to tune the selected node features of the graph structure (selective circuit dimensions) of my two stage CMOS operational amplifier circuit(fixed circuit/graph topology) to get the desigered performance metrics with in its respective predefined ranges and also need to keep the circuit devices(especially MOS transistors) in stable state(saturetion state). I need to know and you need to completely review the below code, whether the implemented computation is properly satisfies my condition or not, in the below code, the state input from the graph is 'node_features_tensor, _, edge_index = state', in that state 'node_features_tensor' has 20 nodes and each node has 24 node features, among the 24 node features of each node, 24th node feature of each node represents its stability condition of the particular node, if it is '1.0000' then the node is stable, else '0.0000' the node is unstable. here the below code we additionally added a heads to make all the nodes in stable state('1.0000'), whether the below implementation satisfies my above said requirements or not? please review and check.

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 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()