Check whether we are accumulating reward of all steps and check whether the policy were designed appropriately for get trained as a good policy for maximizing the accumulated reward.

class PPOAgent:
    def __init__(self, actor_class, critic_class, 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, std=0.0):
        self.actor = actor_class(gnn_model.conv2.out_channels, action_dim, std)  # Initialize actor
        self.critic = critic_class(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()
        self.bounds_low = bounds_low
        self.bounds_high = bounds_high
        self.std = std
        self.epochs = 10  # Define the number of epochs for policy update

    def select_action(self, state, edge_index, edge_attr):
        # Pass state through the GNN model first to get state’s embedding
        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()
    
    def compute_gae(self, next_value, rewards, dones, values, gamma=0.99, lambda_=0.95):
        values.append(next_value)
        gae = 0
        returns = []
        for step in reversed(range(len(rewards))):
            #delta = rewards[step] + gamma * values[step + 1] * dones[step] - values[step]
            delta = rewards[step] + gamma * values[step + 1] * (1 - dones[step]) - values[step]
            #gae = delta + gamma * lambda_ * dones[step] * gae
            gae = delta + gamma * lambda_ * (1 - dones[step]) * gae
            returns.insert(0, gae + values[step])
        return returns
    
    def update_policy(self, states, actions, log_probs, returns, advantages):
        for epoch in range(self.epochs):
            sampler = BatchSampler(SubsetRandomSampler(range(len(states))), batch_size=64, drop_last=True)
            for indices in sampler:
                sampled_states = torch.stack([states[i] for i in indices])
                sampled_actions = torch.stack([actions[i] for i in indices])
                sampled_log_probs = torch.stack([log_probs[i] for i in indices])
                sampled_returns = torch.stack([returns[i] for i in indices])
                sampled_advantages = torch.stack([advantages[i] for i in indices])

                # Assuming the actor model returns a distribution from which log probabilities can be computed
                mean, new_std = self.actor(sampled_states)
                dist = Normal(mean, new_std.exp())
                new_log_probs = dist.log_prob(sampled_actions)
                
                ratio = (new_log_probs - sampled_log_probs).exp()
                
                # Incorporating stability loss
                stability_losses = []
                for state in sampled_states:
                    stability_loss_val = stability_loss(state, target_stability=1.0)  # Assuming you want all nodes to be stable
                    stability_losses.append(stability_loss_val)
                mean_stability_loss = torch.mean(torch.stack(stability_losses))

                surr1 = ratio * sampled_advantages
                surr2 = torch.clamp(ratio, 1.0 - self.epsilon, 1.0 + self.epsilon) * sampled_advantages
                actor_loss = -torch.min(surr1, surr2).mean() + mean_stability_loss
                critic_loss = F.mse_loss(sampled_returns, self.critic(sampled_states))

                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                self.optimizer_actor.step()

                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                self.optimizer_critic.step()


# Training loop
def train(env, agent, num_episodes, max_timesteps, batch_size, epsilon):
    for episode in range(num_episodes):
        node_features_tensor, edge_feature_tensor, edge_index, performance_metrics = env.reset()
        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        dones = []

        state = torch.tensor(node_features_tensor, dtype=torch.float32)  # Assuming the node_features_tensor is a tensor
        for t in range(max_timesteps):
            action, log_prob = agent.select_action((state, edge_feature_tensor, edge_index))
            #next_state, next_edge_feature_tensor, next_edge_index, reward, done, previous_metrics = env.step(action)
            next_state, next_edge_feature_tensor, next_edge_index, reward, done, previous_metrics = env.step(action.numpy())
            print("next_state1", next_state)
            next_state = torch.tensor(next_state, dtype=torch.float32)  # Convert to tensor if not already
            print("next_state2", next_state)
            episode_rewards.append(reward)
            states.append(state)
            actions.append(action)
            log_probs.append(log_prob)
            values.append(agent.critic(state).item())  # Assuming this is how you get value estimation
            dones.append(1 - float(done))
            state = next_state
            edge_feature_tensor = next_edge_feature_tensor
            edge_index = next_edge_index

            if done:
                next_value = agent.critic(next_state).item()  # Fetch next state value for GAE
                break
        
        # Outside the loop, we need to handle the case when we haven’t reached done
        if not done:
            next_value = agent.critic(next_state).item()
        
        # Compute returns and advantages        
        returns = agent.compute_gae(next_value, episode_rewards, dones, values, agent.gamma, agent.lamda)
        
        # Normalizing advantages 
        advantages = torch.tensor(returns) - torch.tensor(values)

        # Update policy and value network
        agent.update_policy(states, actions, log_probs, returns, advantages)

        # Log episode information
        total_reward = sum(episode_rewards)
        print(f"Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}")

# Create the environment
env = CircuitEnvironment(server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content)
gnn_model = EnhancedGNNModelWithSharedParams(num_node_features, num_edge_features, num_out_features)
actor_output_features = gnn_model.conv2.out_channels
print(f"Initializing Actor with output feature size: {actor_output_features}")
agent = PPOAgent(actor, critic, gnn_model, action_dim, bounds_low, bounds_high, lr_actor, lr_critic, gamma, lambda_, epsilon, std)

# Train agent
train(env, agent, env.num_episodes, env.max_timesteps, env.batch_size, env.epsilon)