transistor_regions, saturation_condition, performance_metrics = self.simulator.run_simulation(raw_actions)

the following code 'adjust_raw_actions_value' is my transistor region checking and adjust the dimensions of the transistor if it is not in saturation region. this is the iterative process, we cannot achive the requirement of all transisto to be in saturation with in the single adjustment of the corresponding transistor dimensions which are voilating the required saturation region condition, we need to check the status of the transistor region on each ajustment made in the dimensions by passing it into the simulator 'simulator.run_simulation', with in 2 to 8 loop cycles, the adjustment we made can achive the required criteria of all transistor with in the saturation region.

if all(region == 2 for region in transistor_regions.values()):
    break  # If all transistor regions are '2', exit the loop

# If the condition is not satisfied, adjust the action
raw_actions = adjust_raw_actions_value(raw_actions, transistor_regions)

Implementation Strategy,
# Phase 1: Pre-Training with Manual Tuning Data
1. Generate Dataset: Create a comprehensive dataset of state-action pairs from the manual tuning process.
2. Pre-Train: Pre-train the agent’s neural network on this dataset to ensure it understands the correct tuning dynamics.

# Phase 2: Fine-Tuning with RL Training
3. Hybrid Loss Function: Incorporate both supervised loss and RL loss to refine the agent’s actions.
4. Constrained Action Space: Implement constraints to keep proposed actions within the safe region.
5. Reward Shaping & Penalty: Shape the reward function to provide nuanced feedback and enforce strict penalties for violations.

# Phase 3: Integrated Post-Processing
6. Post-Processing Heuristic: Implement a final manual tuning check after the PPO agent proposes actions. If any transistor is out of saturation, apply manual adjustments before calculating the reward.

for my above requirements, you updated my previous code and you provided the below given updated code with the implementation strategy,

import torch
import torch.optim as optim
import torch.nn.functional as F
from torch_geometric.data import Data
import numpy as np

class CircuitEnvironment:
    def init(self, server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content):
        # Initialization code here
        pass

    def reset(self):
        # Reset environment
        return node_features_tensor, edge_feature_tensor, edge_index, performance_metrics

    def step(self, raw_actions, previous_metrics):
        raw_actions = np.squeeze(raw_actions)
        transistor_regions, saturation_condition, performance_metrics = self.simulator.run_simulation(raw_actions)
        
        # Post-processing step to ensure saturation
        for _ in range(8):  # Limit to 8 iterations
            if all(region == 2 for region in transistor_regions.values()):
                break
            raw_actions = adjust_raw_actions_value(raw_actions, transistor_regions)
            transistor_regions, saturation_condition, performance_metrics = self.simulator.run_simulation(raw_actions)

        node_features_tensor, edge_feature_tensor, edge_index = netlist_to_graph(self.netlist_content, raw_actions, transistor_regions)
        reward = self.calculate_reward(performance_metrics, previous_metrics, transistor_regions, self.previous_transistor_regions)
        done = self.check_done(performance_metrics)
        previous_metrics = performance_metrics
        self.previous_transistor_regions = transistor_regions.copy()

        return node_features_tensor, edge_feature_tensor, edge_index, reward, done, previous_metrics

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

    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)

        return action.detach().numpy(), 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):
        actor_losses = []
        critic_losses = []
        
        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
                action_means, action_stds, state_value = self.actor(node_features_tensor)
                dist = torch.distributions.Normal(action_means, action_stds)
                log_prob = dist.log_prob(action).sum(axis=-1)
                entropy = dist.entropy().sum(axis=-1)
                log_probs_batch.append(log_prob)
                state_values_batch.append(state_value)
                entropies_batch.append(entropy)
                
                # Stability loss
                stability_loss = F.mse_loss(state_value, torch.tensor(returns, dtype=torch.float32))
                stability_losses_batch.append(stability_loss)

            # Convert to tensors
            log_probs_batch = torch.stack(log_probs_batch)
            state_values_batch = torch.stack(state_values_batch)
            entropies_batch = torch.stack(entropies_batch)
            stability_losses_batch = torch.stack(stability_losses_batch)

            # Surrogate function
            ratio = torch.exp(log_probs_batch - prev_log_probs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1.0 - self.policy_clip, 1.0 + self.policy_clip) * advantages
            actor_loss = -torch.min(surr1, surr2).mean() - self.entropy_coef * entropies_batch.mean() + stability_losses_batch.mean()
            critic_loss = F.mse_loss(torch.tensor(returns, dtype=torch.float32), state_values_batch)
            
            self.optimizer_actor.zero_grad()
            actor_loss.backward()
            self.optimizer_actor.step()
            
            self.optimizer_critic.zero_grad()
            critic_loss.backward()
            self.optimizer_critic.step()

            actor_losses.append(actor_loss.item())
            critic_losses.append(critic_loss.item())

        return sum(actor_losses) / len(actor_losses), sum(critic_losses) / len(critic_losses)

def pre_training_dataset_generation():
    """Generate pre-training dataset with manual tuning data"""
    dataset = []
    for _ in range(1000):  # Create 1000 samples
        raw_actions = np.random.uniform(low=env.bounds_low, high=env.bounds_high)  # Random initial actions
        transistor_regions, saturation_condition, performance_metrics = env.simulator.run_simulation(raw_actions)
        
        for _ in range(8):  # Limit to 8 iterations
            if all(region == 2 for region in transistor_regions.values()):
                break
            raw_actions = adjust_raw_actions_value(raw_actions, transistor_regions)
            transistor_regions, saturation_condition, performance_metrics = env.simulator.run_simulation(raw_actions)
        
        node_features_tensor, edge_feature_tensor, edge_index = netlist_to_graph(env.netlist_content, raw_actions, transistor_regions)
        state = (torch.tensor(node_features_tensor, dtype=torch.float32), torch.tensor(edge_feature_tensor, dtype=torch.float32), torch.tensor(edge_index, dtype=torch.long))
        actions = torch.tensor(raw_actions, dtype=torch.float32)
        
        dataset.append((state, actions))
    return dataset

def pre_train_agent(agent, dataset):
    """Pre-train the agent using the dataset"""
    agent.actor.train()
    pretrain_optimizer = optim.Adam(agent.actor.parameters(), lr=1e-4)
    
    for epoch in range(10):  # Train for 10 epochs
        total_loss = 0
        for state, action in dataset:
            pretrain_optimizer.zero_grad()
            action_means, action_stds, _ = agent.actor(state)
            dist = torch.distributions.Normal(action_means, action_stds)
            log_probs = dist.log_prob(action).sum(axis=-1)
            loss = -log_probs.mean()  # Maximize log_probs -> minimize -log_probs
            loss.backward()
            pretrain_optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch}, Loss: {total_loss / len(dataset)}")
    
def train(env, agent, num_episodes, max_timesteps, batch_size):
    dataset = pre_training_dataset_generation()
    pre_train_agent(agent, dataset)

    for episode in range(num_episodes):
        node_features_tensor, edge_feature_tensor, edge_index, performance_metrics = env.reset()
        state = (node_features_tensor, edge_feature_tensor, edge_index)
        
        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        masks = []
        
        for t in range(max_timesteps):
            action, log_prob, perf_metrics = agent.select_action(state, performance_metrics)
            next_node_features_tensor, next_edge_feature_tensor, next_edge_index, reward, done, previous_metrics = env.step(action, perf_metrics)
            next_state = (next_node_features_tensor, next_edge_feature_tensor, next_edge_index)

            states.append(state)
            actions.append(torch.tensor(action))
            log_probs.append(log_prob)
            values.append(agent.critic(state[0]).item())
            
            episode_rewards.append(reward)
            masks.append(1 - done)

            performance_metrics = previous_metrics
            state = next_state

            if done:
                break

        next_value = agent.critic(state[0]).item() if not done else 0
        returns = agent.compute_gae(next_value, episode_rewards, masks, values)
        
        advantages = (torch.tensor(returns) - torch.tensor(values)).detach().numpy()
        returns = torch.tensor(returns).detach().numpy()
        agent.update_policy(states, actions, log_probs, returns, advantages)

# Create the environment
env = CircuitEnvironment(server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content)

# Instantiate the GNN model
gnn_model = CustomGNN(in_channels, out_channels)

# Initialize the PPO Agent
agent = PPOAgent(gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef)

# Train the agent
train(env, agent, num_episodes, max_timesteps, batch_size)