hi, i need an technical advice in my reinforcement learning implementation(PPO agent) with graph neural network, i need to predict the optimal design parameters for my two stage operational amplifier circuit design(choose best channel length, channel width, bias current and compensation capacitor values) from my predefined bounds for each design parameters. each combination parameter set results the performance metrics. here i need to choose optimal design parameters (action space) which gives output results with in my desired range of performance metrics. here my action space may be discrete or continuous(if we decide discrete action space, then we have too large parameter distribution with in my bounds of respective design parameters in the order of total 10^14, so wisely need to choose action space), each design parameter set which gives output results of an set of output performance metrics, optimal values of each design parameters chosen within the bounds is very tedious because, if an response performance metrics may look like progressive if we slightly vary any of the design parameter which results drastic undesired performance metrics, and for some inputs we may achieve particular performance metrics is better other performance metrics may get worse, to satisfy all the metrics is very tedious. I know it is comes in the multi objective optimization process.
Here i am getting the problem, i currently implemented the continuous action space(if you advice me to use discrete action space instead of continuous it it is necessary, that will i accept to modify my existing code), i generate continuous action space, but the result of the generated action space produce the output from my simulation always makes all the transistor region not in saturation, even i implemented an pre training datasets for training the agent to choose the appropriate action space that only need to generate valid action space that results an responce output of the circuit having all the transistor in saturation region. how to strictly introduce the implementation, i apply penalty that deviates the un saturated region and i introduce seperate head in the network which feedback the stability of the node in graph neural network.



def sample_generator(bounds_low, bounds_high, num_samples):
    """Generate samples using Latin Hypercube Sampling."""
    dim = len(bounds_low)
    
    # Create a Latin Hypercube Sampler
    sampler = qmc.LatinHypercube(d=dim)
    
    # Generate normalized samples (0 to 1)
    norm_samples = sampler.random(n=num_samples)
    
    # Scale the samples to the given bounds
    samples = qmc.scale(norm_samples, bounds_low, bounds_high)
    
    return samples

def pre_training_dataset_generation(env, required_samples=1000, progress_filename='progress.pkl'):
    """Generate pre-training dataset with manual tuning data"""
    env.connect_to_simulator() # Ensure simulator is connected before dataset generation

    dataset, start_sample, raw_samples, saved_required_samples = load_progress(progress_filename)
    print("resume from the sample", start_sample)
    # Check if required_samples is updated, if so, regenerate samples
    if saved_required_samples != required_samples:
        # Generate samples using LHS
        raw_samples = sample_generator(env.bounds_low, env.bounds_high, required_samples)
        start_sample = 0  # restart if required_samples changed
        dataset = []
    # Initialize or append CSV file
    csv_file_mode = 'a' if start_sample > 0 else 'w'
    with open('transistor_data.csv', mode=csv_file_mode, newline='') as csvfile:
        fieldnames = ['L1', 'L3', 'L5', 'L6', 'L7', 'W1', 'W3', 'W5', 'W6', 'W7',
                      'Ib', 'Cc', 'Vc', 'M1', 'M2', 'M3', 'M4', 'M5', 'M8', 'M6', 
                      'M7', 'Area', 'PowerDissipation', 'SlewRate', 'Gain', 'Bandwidth3dB', 
                      'UnityGainFreq', 'PhaseMargin']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        if start_sample == 0:
            writer.writeheader()

        current_sample = start_sample
        while current_sample < required_samples:
            sample_idx = current_sample % len(raw_samples)
            raw_actions = raw_samples[sample_idx]
            valid_sample_found = False

            while not valid_sample_found:
                iteration = 0

                while iteration < env.max_iterations:
                    #print("raw_actions", raw_actions)
                    transistor_regions, saturation_condition, performance_metrics = env.simulator.run_simulation(raw_actions)
                    # Write results directly to CSV
                    writer.writerow({
                        'L1': raw_actions[0], 'L3': raw_actions[1], 'L5': raw_actions[2], 'L6': raw_actions[3], 'L7': raw_actions[4],
                        'W1': raw_actions[5], 'W3': raw_actions[6], 'W5': raw_actions[7], 'W6': raw_actions[8], 'W7': raw_actions[9],
                        'Ib': raw_actions[10], 'Cc': raw_actions[11], 'Vc': raw_actions[12],
                        'M1': transistor_regions.get('M1', None), 'M2': transistor_regions.get('M2', None), 'M3': transistor_regions.get('M3', None),
                        'M4': transistor_regions.get('M4', None), 'M5': transistor_regions.get('M5', None), 'M8': transistor_regions.get('M8', None),
                        'M6': transistor_regions.get('M6', None), 'M7': transistor_regions.get('M7', None),
                        'Area': performance_metrics.get('Area', None), 'PowerDissipation': performance_metrics.get('PowerDissipation', None),
                        'SlewRate': performance_metrics.get('SlewRate', None), 'Gain': performance_metrics.get('Gain', None),
                        'Bandwidth3dB': performance_metrics.get('Bandwidth3dB', None), 'UnityGainFreq': performance_metrics.get('UnityGainFreq', None),
                        'PhaseMargin': performance_metrics.get('PhaseMargin', None)
                    })

                    if all(region == 2 for region in transistor_regions.values()): 
                        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))
                        valid_sample_found = True
                        break
                    # Try adjusting raw actions and continue iterations
                    raw_actions = env.adjust_raw_actions_value(raw_actions, transistor_regions)
                    print("transistor_regions", transistor_regions)
                    print("iteration count", iteration)
                    iteration += 1
                    
                if not valid_sample_found:
                    new_raw_action = sample_generator(env.bounds_low, env.bounds_high, 1)
                    raw_actions = new_raw_action.flatten()

            save_progress(dataset, current_sample, raw_samples, required_samples, progress_filename)
            print("sample count", current_sample)
            current_sample += 1

    env.disconnect_from_simulator() # Optionally disconnect after dataset generation

    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)
    scheduler = torch.optim.lr_scheduler.StepLR(pretrain_optimizer, step_size=30, gamma=0.1)

    for epoch in range(100): # 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()
            
            # Gradient Clipping
            torch.nn.utils.clip_grad_norm(agent.actor.parameters(), max_norm=1.0)
            pretrain_optimizer.step()
            
            total_loss += loss.item()
        scheduler.step()
        print(f"Epoch {epoch}, Loss: {total_loss / len(dataset)}")


def train(env, agent, num_episodes, max_timesteps, batch_size):
    # Ensure simulator connection
    env.connect_to_simulator()
    
    total_samples=1000
    progress_filename='progress.pkl'
    
    start_episode = 0    
    rewards_log = []
    losses_log = []
    metrics_log = []
    episode_lengths = []
    
    dataset = None
    pretrain_params = None
    
    # Check if a previous training state exists and load it
    saved_state = load_training_state()
    if saved_state:
        start_episode = saved_state['episode']
        agent.actor.load_state_dict(saved_state['actor_state_dict'])
        agent.critic.load_state_dict(saved_state['critic_state_dict'])
        agent.optimizer_actor.load_state_dict(saved_state['optimizer_actor_state_dict'])
        agent.optimizer_critic.load_state_dict(saved_state['optimizer_critic_state_dict'])
        rewards_log = saved_state['rewards_log']
        losses_log = saved_state['losses_log']
        print(f"Resuming from episode {start_episode}")
    
    # Check if a previous pre-training state (dataset) exists and load it
    pretraining_state = load_pretraining_state()
    if pretraining_state:
        dataset = pretraining_state['dataset']
        pretrain_params = pretraining_state['pretrain_params']
        print("Resuming pre-training from saved state.")

    # If no pretraining state found, generate dataset and pre-train agent
    if not dataset:
        pretrain_params = {
            'bounds_low': env.bounds_low,
            'bounds_high': env.bounds_high,
            # Add other relevant parameters here
        }
        dataset = pre_training_dataset_generation(env, total_samples, progress_filename)
        pre_train_agent(agent, dataset)
        save_pretraining_state(dataset, pretrain_params)
    else:
        # Pre-training using existing dataset
        pre_train_agent(agent, dataset)
       
    for episode in range(start_episode, 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 = []
        perf_metrics_log = []
    
        for t in range(max_timesteps):
            # Selecting action based on the current state
            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)