I need you to check the computation of actor method in the Class Actor and its followed computation in the 'select_action' method and the computation of actor method in the Class Actor and the follow up computation of 'evaluate' in the 'update_policy' with in the class PPOAgent. its seems both the computation which is with in the select action and the evaluate are getting differs, is it acceptable to implement or it does need to change? if it need to change please provide me the updated code for the same.

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)


    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]]
        
        print("selected_features", selected_features)
        
        # 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)

        return action_means, action_std

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)
        returns = torch.tensor(returns)
        prev_log_probs = torch.tensor(prev_log_probs)
        
        # To extract the 24th feature of each node which indicates stability
        stabilities = [state[0][:, 23] for state in prev_states]
        # Convert list to tensor
        stability_tensor = torch.stack(stabilities)
        
        stability_loss = self.compute_stability_loss(stability_tensor, target_stability=1.0)
        
        for _ in range(self.epochs):
            log_probs, state_values, entropy = self.evaluate(prev_states, prev_actions)             
            
            ratios = torch.exp(log_probs - prev_log_probs.detach())
            
            advantages = 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() + stability_loss
            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()

    def evaluate(self, states, actions):
        action_probs, state_values = [], []
        log_probs, entropy = [], []
        
        for state, action in zip(states, actions):
            # Obtain the model’s action predictions for the given state
            prob, _ = self.actor(state) # prob should preferably be in the form [mean action values]

            value = self.critic(state[0])  
            
            # Compute the variance of predicted actions and ensure a minimal variance to avoid degenerate distributions
            action_variance = prob.var(0) + 1e-5  # Adding a small epsilon for numerical stability
            
            # The model predicts 13 distinct actions, create a tensor of variances for each action, and we want to maintain the same variance across all actions
            variances = action_variance.repeat(13)  # Replace 13 with the dynamic size of prob if necessary
            
            # Construct the covariance matrix
            cov_mat = torch.diag(variances)            
           
            # Define the multivariate normal distribution using the predicted actions (prob) and the covariance matrix
            dist = MultivariateNormal(prob, cov_mat)            
            
            # Ensure ‘action’ is a tensor. Adjust dtype and device as necessary.
            if not isinstance(action, torch.Tensor):
                action_tensor = torch.tensor(action, dtype=prob.dtype, device=prob.device)
            else:
                action_tensor = action

            # Compute log probabilities or sample actions here based on the distribution
            log_prob = dist.log_prob(action_tensor)
            ent = dist.entropy()
            
            # Collect the computed values
            action_probs.append(prob)
            state_values.append(value)
            log_probs.append(log_prob)
            entropy.append(ent)
        
        # Concatenate lists into tensors for batch processing
        action_probs = torch.stack(action_probs)
        state_values = torch.stack(state_values).squeeze(-1)
        log_probs = torch.stack(log_probs)
        entropy = torch.stack(entropy)         
        
        return log_probs, state_values, entropy

    def compute_stability_loss(self, stabilities, target_stability=1.0):
        """Compute stability loss based on stabilities tensor."""
        stability_loss = F.binary_cross_entropy_with_logits(stabilities, torch.full_like(stabilities, fill_value=target_stability))
        return stability_loss