In my reinforcement learning algorithm implementation, i need to pass the intermediate reward to the agent, wen we get 'reward = -200' or 'reward >= 90'. 

In the below existing code i have implemented the reward function, but the implemented reward is only given to the agent after the maximum step size of that current episode, but i need to give intermediate reward with in the episode to guide the agent more effectively. So, i updated the existing code based on my logic, but i dont know whether the implemented logic to pass the intermediate reward in to the PPO agent is correct or wrong. i need you to check completely the logic and procedure for passing the intermediate reward in to the PPO agent is correct or wrong in the updated code. advice me if any modification were needed in the updated code. (or do we required to modify the compute_gae and update_policy function present in the class PPO agent)

Existing Code:
# Define Circuit Environment
class CircuitEnvironment:
    VERY_CLOSE_REWARD = 40
    LARGE_REWARD = 10
    VERY_HIGH_REWARD = 50
    PENALTY = 20
    SEVERE_PENALTY = 200

    def __init__(self, bounds_low, bounds_high):
        # Initialization code: setup bounds
    
    def reset(self, target_metrics=None):
	#existing implementation
        return node_features_tensor, edge_index, performance_metrics    

    def step(self, raw_actions, previous_metrics, target_metrics):
        raw_actions = np.squeeze(raw_actions)
        performance_metrics, temp_netlist_path = run_simulation(raw_actions)
        node_features_tensor, edge_index = netlist_to_graph(temp_netlist_path)
        self.target_metrics = target_metrics 
        
        # Round off the previous and current performance metrics to integer values for comparison
        rounded_performance_metrics = self.round_to_discrete_values(performance_metrics)
        rounded_previous_metrics = self.round_to_discrete_values(previous_metrics)
        print("performance_metrics", rounded_performance_metrics)
        
        reward = self.calculate_reward(rounded_performance_metrics, rounded_previous_metrics)
        done = self.check_done(rounded_performance_metrics)
        # Update previous state
        previous_metrics = rounded_performance_metrics 
        self.node_features_tensor = node_features_tensor
        self.edge_index = edge_index
        return node_features_tensor, edge_index, reward, done, previous_metrics

    def calculate_reward(self, c_metrics, p_metrics):
        metrics_order = ['gain', 'bw', 'ugf', 'pm']
            
        reward = 0

        target_metrics = {k: self.target_metrics[i] for i, k in enumerate(metrics_order)}
        print("target_metrics", target_metrics)
        for metric_name in metrics_order:
            c_value = c_metrics[metric_name]
            p_value = p_metrics[metric_name]
            target_value = target_metrics[metric_name]
            
            if c_value == target_value:
                reward += self.VERY_HIGH_REWARD
            elif abs(c_value - target_value) < abs(p_value - target_value):
                reward += self.LARGE_REWARD
            else:
                reward -= self.PENALTY
                
        return reward
        
    def check_done(self, performance_metrics):
        
        if self.target_metrics is None:
            return False
        metrics_order = ['gain', 'bw', 'ugf', 'pm']
        target_metrics = {k: self.target_metrics[i] for i, k in enumerate(metrics_order)}
        for metric_name in metrics_order:
            if performance_metrics[metric_name] != target_metrics[metric_name]:
                return False
        return True

# Training Function
def train(env, agent, num_episodes, max_timesteps):
    num_targets = 100
    train_targets = CircuitEnvironment.generate_random_targets(num_targets=num_targets, seed=0)
    test_targets = CircuitEnvironment.generate_random_targets(num_targets=20, seed=1)
    
    # Initialize the episode to start from and log lists
    start_episode = 0    
    rewards_log = []
    losses_log = []
    metrics_log = []
    episode_lengths = []
    
    # Attempt to load the saved state
    saved_state = load_training_state()
        
    for episode in range(start_episode, num_episodes):
        if episode >= num_targets:
            break  # Exit the loop if exhausted the number of targets
        
        target = train_targets[episode % len(train_targets)]
        
        node_features_tensor, edge_index, performance_metrics = env.reset(target)
        state = (node_features_tensor, edge_index)
    
        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        masks = []
        perf_metrics_log = []

        for t in range(max_timesteps):            
            action, log_prob, perf_metrics = agent.select_action(state, performance_metrics)
            next_node_features_tensor, next_edge_index, reward, done, previous_metrics = env.step(action, perf_metrics, target)
            
            next_state = (next_node_features_tensor, next_edge_index)
    
            episode_rewards.append(reward)
            states.append(state)
            actions.append(action)
            log_probs.append(log_prob.item())            
            values.append(agent.critic(node_features_tensor, edge_index).item())
            masks.append(1.0 - float(done))
            perf_metrics_log.append(perf_metrics)
                
            performance_metrics = previous_metrics
            state = next_state            
  
            if done:
                break
    
        next_value = agent.critic(state[0], state[1]).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()
        average_loss_actor, average_loss_critic = agent.update_policy(states, actions, log_probs, returns, advantages)
    
        agent.save_checkpoint(episode)        

        save_training_state(agent, episode, rewards_log, losses_log, metrics_log, train_targets, test_targets)
    
    agent.save_checkpoint('final')        
   
if __name__ == "__main__":
    # Initialize environment parameters
    # Define the values for bounds, and target metrics

    env = CircuitEnvironment(bounds_low, bounds_high)
    agent = PPOAgent(state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef)

    train(env, agent, num_episodes, max_timesteps)


Updated Code:
def train(env, agent, num_episodes, max_timesteps):
    num_targets = 100
    train_targets = CircuitEnvironment.generate_random_targets(num_targets=num_targets, seed=0)
    test_targets = CircuitEnvironment.generate_random_targets(num_targets=20, seed=1)
    
    # Initialize the episode to start from and log lists
    start_episode = 0    
    rewards_log = []
    losses_log = []
    metrics_log = []
    episode_lengths = []
    
    # Attempt to load the saved state
    saved_state = load_training_state()
        
    for episode in range(start_episode, num_episodes):
        if episode >= num_targets:
            break  # Exit the loop if exhausted the number of targets
        
        target = train_targets[episode % len(train_targets)]
        
        node_features_tensor, edge_index, performance_metrics = env.reset(target)
        state = (node_features_tensor, edge_index)
    
        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        masks = []
        perf_metrics_log = []

        for t in range(max_timesteps):            
            action, log_prob, perf_metrics = agent.select_action(state, performance_metrics)
            next_node_features_tensor, next_edge_index, reward, done, previous_metrics = env.step(action, perf_metrics, target)
            
            next_state = (next_node_features_tensor, next_edge_index)
    
            episode_rewards.append(reward)
            states.append(state)
            actions.append(action)
            log_probs.append(log_prob.item())            
            values.append(agent.critic(node_features_tensor, edge_index).item())
            masks.append(1.0 - float(done))
            perf_metrics_log.append(perf_metrics)
                
            performance_metrics = previous_metrics
            state = next_state
            
            # Check for intermediate reward conditions
            if reward <= -200 or reward >= 90:
                next_value = agent.critic(state[0], state[1]).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()
                average_loss_actor, average_loss_critic = agent.update_policy(states, actions, log_probs, returns, advantages)
                
                # Reset episodic storage after intermediate update
                episode_rewards = []
                states = []
                actions = []
                log_probs = []
                values = []
                masks = []

            if done:
                break
    
        # Final policy update at the end of the episode
        next_value = agent.critic(state[0], state[1]).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()
        average_loss_actor, average_loss_critic = agent.update_policy(states, actions, log_probs, returns, advantages)
        
        agent.save_checkpoint(episode)        
        save_training_state(agent, episode, rewards_log, losses_log, metrics_log, train_targets, test_targets)
    
    agent.save_checkpoint('final')