I need to implement the RL with GNN algorithm to optimize my analog circuit design parameters and need to get the output performance metrics(with in the target specification range), I need to tune the design parameters with in the predefined bounds of each design variables to achieve my objective. And an important note, you should know that my circuit schematic/circuit netlist is fixed layout, only the design parameters are need to tune with in the bounds to get my target objective.

I converted my circuit netlist in to undirected bipartite graph (with circuit components and associated nets as nodes) each edges are connected between the component nodes(11 component nodes) and net nodes(9 net nodes), in such a way that the graph are arranges by component nodes on one side and the net nodes on other side. here the graph layout is fixed i only need to tune the node features of component nodes alone, i need not to care about the net nodes, because it is fixed/ideal throughout the process the net features wont get any change. i need to concentrate on only node features, especially the feature ‘values’ belongs to each component nodes, it needs to be tune in 'continuous action space' with in its respective bounds. I need to tune only the component nodes, and especially the feature ‘values’ belongs to each component nodes.

My existing 'node_features_tensor' are taken as per the below given description:
node_features_tensor = device_type + device_onehot + component_onehot + values + region_state, with indices [0:23] which is splitted for each features as, [0 + 1:6 + 7:17 + 18:22 + 23]
we can differentiate the component node and the net node from the feature 'device_type' at the index [0] if the value is '0.0000' then it is component node if the value is '1.0000' then it is net node(ensure the features for which remain unchange)

the indices of node_features_tensor [7:17] represents the component index number to identify the components  (‘M0’ to ‘M7’, ‘C0’, ‘I0’, ‘V1’) this attribute is used to identify the component name within the component node. which is one-hot state representation of corresponding component.

the indices of node_features_tensor [18:22] represents the values of coresponding components 'w_value', 'l_value', 'C_value', 'I_value', 'V_value', these are the parameter need to be optimize by the algorithm. 

we have totally 11 components (‘M0’ to ‘M7’, ‘C0’, ‘I0’, ‘V1’), in that ‘M0’ to ‘M7’ components has two value features to tune 'w_value', 'l_value', and ‘C0’ component has one value feature to tune 'C_value', and ‘I0’ component has one value feature to tune 'I_value',  and ‘V1’ component has one value feature to tune 'V_value'
the indices [18:19] are the normalized values of the 'w_value', 'l_value' correspondind to the components ‘M0’ to ‘M7’,
the index [20] are the normalized value of the 'C_value' correspondind to the component ‘C0’,
the index [21] are the normalized value of the 'I_value' correspondind to the component ‘I0’,
the index [22] are the normalized value of the 'V_value' correspondind to the component ‘V1’,

What is the correct way of representing features to the GNN model by preprocessing before training, with considering particular required components node only need to be tune during the training, and for that particular component only the the particular values of the corresponding node to be tunes during training remains unnchange, for sample among all nodes we need to consider only the component nodes to tune by selecting the component nodes from the node_features_tensor, 
To select the component nodes: 
from the node_features_tensor, index of [0] with the value '0.0000' to be select, 
Then, to select each component node: 
from the node_features_tensor, indices of [7:17]
for 'M0' index [7] has value '1.0000', 'M1' index [8] has value '1.0000', 'M2' index [9] has value '1.0000', 'M3' index [10] has value '1.0000', 'M4' index [11] has value '1.0000', 'M5' index [12] has value '1.0000', 'M6' index [13] has value '1.0000', 'M7' index [14] has value '1.0000', 'C0' index [15] has value '1.0000', 'I0' index [16] has value '1.0000', 'V1' index [17] has value '1.0000'.
Then we need to tune only the corresponding values of each components during training:
we need to focus only the particular values index for particular component node index,
from the node_features_tensor, indices of [18:22]
for the component nodes  'M0' index [7], 'M1' index [8], 'M2' index [9], 'M3' index [10], 'M4' index [11], 'M5' index [12], 'M6' index [13], 'M7' index [14] we need to focus only the  node_features_tensor, indices of [18:19] for tuning its 'values' during the training process,
for the component node  'C0' index [15], we need to focus only the  node_features_tensor, index of [20] for tuning its 'values' during  the training process,
for the component node  'I0' index [16], we need to focus only the  node_features_tensor, index of [21] for tuning its 'values' during  the training process
for the component node  'V1' index [17], we need to focus only the  node_features_tensor, index of [22] for tuning its 'values' during  the training process

My existing implementation:

import torch
from torch_geometric.nn import GATConv, NNConv
import torch.nn.functional as F
from torch.nn import Sequential, Linear, ReLU

class EnhancedGNNModelWithSharedParams(torch.nn.Module):
    def init(self, num_node_features, num_edge_features, num_out_features, shared_param_dims):
        super(EnhancedGNNModelWithSharedParams, self).init()
        self.conv1 = GATConv(num_node_features, 8, heads=8, dropout=0.6)
        
        # Define shared parameters for paired components
        self.shared_params = torch.nn.ParameterDict({
            ‘M0_M1’: torch.nn.Parameter(torch.randn(shared_param_dims)),
            ‘M2_M3’: torch.nn.Parameter(torch.randn(shared_param_dims)),
            ‘M4_M7’: torch.nn.Parameter(torch.randn(shared_param_dims))
        })

        nn_model = Sequential(Linear(num_edge_features, 32), ReLU(), Linear(32, 8 * 8))
        self.conv2 = NNConv(8 * 8, num_out_features, nn_model, aggr=‘mean’)

        # Placeholder for gradient_mask, will be defined in replace_with_shared_params
        self.gradient_mask = None

    def forward(self, node_features, edge_index, edge_attr):
        # Apply shared params modification
        modified_node_features = self.replace_with_shared_params(node_features)
        
        # Before convolutions, apply gradient mask
        if self.gradient_mask is not None:
            modified_node_features.register_hook(lambda grad: grad * self.gradient_mask)
        
        x = F.dropout(modified_node_features, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index, edge_attr)
        return x
    
    def replace_with_shared_params(self, node_features):
        device_types = node_features[:, 0]  # Assuming the first column is device_type
        component_mask = (device_types == 0.0000)
        
        # Define value indices and the pairs
        value_indices = [18, 19, 20, 21, 22]
        pairs = [(‘M0’, ‘M1’), (‘M2’, ‘M3’), (‘M4’, ‘M7’)]
        
        # Initialize gradient mask
        self.gradient_mask = torch.zeros_like(node_features, dtype=torch.float32)
        
        for pair in pairs:
            shared_param = self.shared_params[f"{pair[0]}_{pair[1]}"]
            for comp in pair:
                # Direct assignment based on predefined mappings
                component_index = identification_indices[comp]
                node_features[component_mask, value_indices] = shared_param

        # Update gradient mask for component value indices
        self.gradient_mask[component_mask, value_indices] = 1.0
        
        return node_features

def rearrange_action_output(raw_action_tensor):
    # Ensure raw_action_tensor is in the correct shape, e.g., [13]
    # Rearrange indices to match the specified order
    rearranged_indices = [1, 2, 4, 6, 7, 8, 9, 0, 3, 5, 11, 10, 12]  # Corresponds to [‘l1_value’, ‘l3_value’, ‘l5_value’, ‘w6_value’, ‘l6_value’, ‘w7_value’, ‘l7_value’, ‘w1_value’, ‘w3_value’, ‘w5_value’, ‘Ib_value’, ‘Cc_value’, ‘Vc_value’]
    return raw_action_tensor[rearranged_indices]

class PPOAgent:
    def init(self, actor, critic, gnn_model, action_dim, bounds_low, bounds_high,
                 lr_actor=3e-4, lr_critic=1e-3, gamma=0.99, lamda=0.95, epsilon=0.2):
        self.actor = actor(gnn_model.conv2.out_channels, action_dim)  # Initialize actor
        self.critic = critic(gnn_model.conv2.out_channels)  # Initialize critic
        self.gnn_model = gnn_model  # GNN model instance
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr_critic)
        self.gamma = gamma
        self.lamda = lamda
        self.epsilon = epsilon
        self.bounds_low = torch.tensor(bounds_low).float()
        self.bounds_high = torch.tensor(bounds_high).float()

    def select_action(self, state, edge_index, edge_attr):
        # Pass state through the GNN model first to get state’s embedding
        with torch.no_grad():
            state_embedding = self.gnn_model(state, edge_index, edge_attr)
        
        # Then, pass the state embedding through the actor network to get action mean and std
        mean, std = self.actor(state_embedding)

        # Rearranging based on the specifications
        rearranged_mean = rearrange_action_output(mean)

        # Scale mean based on action bounds defined
        mean = self.bounds_low + (torch.sigmoid(rearranged_mean) * (self.bounds_high - self.bounds_low))
        dist = Normal(mean, std)
        
        # Sample an action from the distribution and calculate its log probability
        action = dist.sample()

        action = torch.clamp(action, self.bounds_low, self.bounds_high)  # Ensure action is within bounds
        action_log_prob = dist.log_prob(action)
        
        return action.detach(), action_log_prob.detach()

gnn_model = EnhancedGNNModelWithSharedParams(num_node_features, num_edge_features, num_out_features, shared_param_dims)
ppo_agent = PPOAgent(actor, critic, gnn_model, action_dim, lr_actor, lr_critic, gamma, lambda_, epsilon)