class Actor(torch.nn.Module):
    def __init__(self, gnn_model):
        super(Actor, self).__init__()
        self.gnn = gnn_model

    def forward(self, state):
        # State contains node_features_tensor, edge_feature_tensor, edge_index
        node_features_tensor, edge_feature_tensor, edge_index = state
        action_probs = self.gnn(node_features_tensor, edge_index)
        return action_probs

class Critic(torch.nn.Module):
    def __init__(self, state_dim):
        super(Critic, self).__init__()
        self.network = torch.nn.Sequential(
            torch.nn.Linear(state_dim, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 1)
        )

    def forward(self, state):
        return self.network(state)

class PPOAgent:
    def init(self, gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs):
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.epsilon = epsilon
        self.policy_clip = policy_clip
        self.epochs = epochs
        
        self.actor = Actor(gnn_model)
        self.critic = Critic(state_dim)
        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  # Assume continuous

    def select_action(self, state):
        state_tensor = torch.FloatTensor(state).unsqueeze(0)  # Adjust dimensions as necessary
        action_probs = self.actor(state_tensor)
        cov_mat = torch.diag(action_probs.var()).unsqueeze(0)  # Ensure variances are positive and form a covariance matrix
        dist = MultivariateNormal(action_probs, cov_mat)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        
        return action.numpy().squeeze(), log_prob.item()
    
    def compute_gae(self, next_value, rewards, masks, values):
        values = values + [next_value]
        gae = 0
        returns = []
        for step in reversed(range(len(rewards))):
            delta = rewards[step] + self.gamma * values[step + 1] * masks[step] - values[step]
            gae = delta + self.gamma * self.gae_lambda * masks[step] * gae
            returns.insert(0, gae + values[step])
        return returns

    def update_policy(self, prev_states, prev_actions, prev_log_probs, returns, advantages):
        advantages = torch.tensor(advantages)
        returns = torch.tensor(returns)
        prev_log_probs = torch.tensor(prev_log_probs)
        
        for _ in range(self.epochs):
            log_probs, state_values, entropy = self.evaluate(prev_states, prev_actions)
            ratios = torch.exp(log_probs - prev_log_probs.detach())
            
            advantages = returns - state_values.detach()
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1-self.policy_clip, 1+self.policy_clip) * advantages
            
            actor_loss = - torch.min(surr1, surr2).mean()
            critic_loss = F.mse_loss(state_values, returns)
            
            self.optimizer_actor.zero_grad()
            actor_loss.backward()
            self.optimizer_actor.step()
            
            self.optimizer_critic.zero_grad()
            critic_loss.backward()
            self.optimizer_critic.step()

    def evaluate(self, states, actions):
        # Replace with actual evaluation logic based on your training loop requirements
        pass

Check the below training loop is properly synchronized with the above given 'class PPOAgent' and all its functions 'select_action', 'compute_gae', 'update_policy' exactly,

# 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()
        state = (node_features_tensor, edge_index)  # Adjust based on your actual state representation needs
        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        masks = []

        for t in range(max_timesteps):
            action, log_prob = agent.select_action(state)
            next_state, next_edge_feature_tensor, next_edge_index, reward, done, previous_metrics = env.step(action.numpy())
            next_state = torch.tensor(next_state, dtype=torch.float32)  # Convert to tensor if not already
            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
            masks.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)
        
        # 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}")