I need your review in the assigned initial training parameters(hyper parameters values), I need you to make it as optimal value, here I used epochs value as 10 is it sufficient, in my training process, each training phase is processed upto maximum 300 steps to reach the target.


# PPOAgent definition
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
        
        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 
        
        self.log_dir = 'OPAMP_GNNRL_LTS_TrainTest5_log_file'
        os.makedirs(self.log_dir, exist_ok=True)
        
        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 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)

        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.bounds_low + (self.bounds_high - self.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 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 Vc directly
        rounded_actions[7] = np.round(scaled_actions_np[7], 1)  # for Vc

        return rounded_actions, log_probs.detach(), performance_metrics

    def compute_gae(self, next_value, rewards, masks, values):
        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  
        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)
                
                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)
             
            # 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) 
            entropy = torch.cat(entropies_batch)
            
            # Surrogate function
            ratios = torch.exp(log_probs - prev_log_probs)            
            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(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())  
            critic_losses.append(critic_loss.item())
            
        # 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

if __name__ == "__main__":
    # Initialize environment parameters
    # Define the values for bounds, and target metrics
    bounds_low = np.array([0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 15e-6, 0.1e-12, 0.8])
    bounds_high = np.array([50e-6, 50e-6, 50e-6, 50e-6, 50e-6, 30e-6, 10e-12, 1.4])
    # Initialize Hyperparameters
    state_dim = 23
    action_space = 8
    lr_actor = 1e-4 # 1e-4
    lr_critic = 2e-4 # 3e-4
    gamma = 0.99
    gae_lambda = 0.95
    epsilon = 0.2
    policy_clip = 0.2
    epochs = 10
    entropy_coef = 0.01

    env = CircuitEnvironment(bounds_low, bounds_high)
    agent = PPOAgent(state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef)

    num_episodes = 1000
    max_timesteps = 300

    train(env, agent, num_episodes, max_timesteps)

    trained_state = load_training_state()
    if trained_state:
        agent.actor.load_state_dict(trained_state['actor_state_dict'])
        agent.critic.load_state_dict(trained_state['critic_state_dict'])

    num_test_targets = 20
    test(env, agent, num_test_targets, max_timesteps_per_target=300)