in this below code, all the necessary datas are saved in the training process within the function 'train' and the saved state were loaded whenever we needed pause and play the iteration process, similarly, i need to save load the test process with in the function 'test'. so i need you to update the code to save the necessary all needed data in the test process to save and load.

# PPOAgent definition
class PPOAgent:
    def __init__(self, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef):     
        self.actor = Actor(gnn_model_actor)
        self.critic = Critic(gnn_model_critic)
        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 
        
        self.log_dir = 'OPAMP_GNNRL_LTS_TrainTest3_log_file'
        os.makedirs(self.log_dir, exist_ok=True)
        
    def save_checkpoint(self, episode):
        torch.save(self.actor.state_dict(), os.path.join(self.log_dir, f'actor_{episode}.pt'))
        torch.save(self.critic.state_dict(), os.path.join(self.log_dir, f'critic_{episode}.pt'))

    def load_checkpoint(self, episode):
        self.actor.load_state_dict(torch.load(os.path.join(self.log_dir, f'actor_{episode}.pt')))
        self.critic.load_state_dict(torch.load(os.path.join(self.log_dir, f'critic_{episode}.pt')))
              
# save and load process 
def save_training_state(agent, episode, rewards_log, losses_log, metrics_log, train_targets, test_targets, filepath='training_state_TrainTest3.pt'):
    state = {
        'episode': episode,
        'actor_state_dict': agent.actor.state_dict(),
        'critic_state_dict': agent.critic.state_dict(),
        'optimizer_actor_state_dict': agent.optimizer_actor.state_dict(),
        'optimizer_critic_state_dict': agent.optimizer_critic.state_dict(),
        'rewards_log': rewards_log,
        'losses_log': losses_log,
        'metrics_log': metrics_log,
        'train_targets': train_targets,
        'test_targets': test_targets
    }
    torch.save(state, filepath)

def load_training_state(filepath='training_state_TrainTest3.pt'):
    if os.path.isfile(filepath):
        state = torch.load(filepath)
        return state
    else:
        return None

# Training Function
def train(env, agent, num_episodes, max_timesteps):
    train_targets = CircuitEnvironment.generate_random_targets(num_targets=100, 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()
    if saved_state:
        start_episode = saved_state['episode'] + 1  # Start from the next 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']
        metrics_log = saved_state['metrics_log']
        train_targets = saved_state['train_targets']
        test_targets = saved_state['test_targets']
        print(f"Resuming from episode {start_episode}")
        
    for episode in range(start_episode, num_episodes): 
        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)                        
            
        episode_length = len(episode_rewards)
        total_reward = sum(episode_rewards)
        mean_episode_reward = total_reward / episode_length
        deployment_accuracy = calculate_accuracy(perf_metrics_log, target)
        
        min_episode_reward = min(episode_rewards)
        max_episode_reward = max(episode_rewards)
        
        # Metrics calculations
        metrics_agg = {}
        for metric in ['gain', 'bw', 'ugf', 'pm']:
            metric_values = [m[metric] for m in perf_metrics_log]
            metrics_agg[f'{metric}_min'] = min(metric_values)
            metrics_agg[f'{metric}_max'] = max(metric_values)
            metrics_agg[f'{metric}_mean'] = sum(metric_values) / len(metric_values)
        
        # Log episode information            
        print(f'Train: Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}, Mean Episode Reward: {mean_episode_reward}, Min Episode Reward: {min_episode_reward}, Max Episode Reward: {max_episode_reward}, Deployment Accuracy: {deployment_accuracy}')
             
        # Logging
        logging.info(f" Train: Episode {episode + 1} - Reward: {total_reward}, Length: {episode_length}, Actor Loss: {average_loss_actor}, Critic Loss: {average_loss_critic}")
        logging.info(f"Train:  Episode {episode + 1} - Mean Episode Reward: {mean_episode_reward}, Min Episode Reward: {min_episode_reward}, Max Episode Reward: {max_episode_reward}, Deployment Accuracy: {deployment_accuracy}")   
        logging.info(f"Train: Episode {episode + 1} - Performance Metrics (min, max, mean): {metrics_agg}")

        rewards_log.append({'episode': episode + 1, 'total_reward': total_reward})
        losses_log.append({'episode': episode + 1, 'actor_loss': average_loss_actor, 'critic_loss': average_loss_critic})
        metrics_log.append({'episode': episode + 1, 'Mean Episode Reward': mean_episode_reward, 'Min Episode Reward': min_episode_reward, 'Max Episode Reward': max_episode_reward, 'Deployment Accuracy': deployment_accuracy, **metrics_agg})
    
        episode_lengths.append(episode_length)
    
        agent.save_checkpoint(episode)
        
        pd.DataFrame(rewards_log).to_csv(os.path.join(agent.log_dir, 'rewards_log.csv'), index=False)
        pd.DataFrame(losses_log).to_csv(os.path.join(agent.log_dir, 'losses_log.csv'), index=False)
        pd.DataFrame(metrics_log).to_csv(os.path.join(agent.log_dir, 'metrics_log.csv'), index=False)
        save_training_state(agent, episode, rewards_log, losses_log, metrics_log, train_targets, test_targets)
        '''
        if episode % 5 == 0 or episode == num_episodes - 1:
            pd.DataFrame(rewards_log).to_csv(os.path.join(agent.log_dir, 'rewards_log.csv'), index=False)
            pd.DataFrame(losses_log).to_csv(os.path.join(agent.log_dir, 'losses_log.csv'), index=False)
            pd.DataFrame(metrics_log).to_csv(os.path.join(agent.log_dir, 'metrics_log.csv'), index=False)
            save_training_state(agent, episode, rewards_log, losses_log, metrics_log, train_targets, test_targets)
        '''
    pd.DataFrame(rewards_log).to_csv(os.path.join(agent.log_dir, 'final_rewards_log.csv'), index=False)
    pd.DataFrame(losses_log).to_csv(os.path.join(agent.log_dir, 'final_losses_log.csv'), index=False)
    pd.DataFrame(metrics_log).to_csv(os.path.join(agent.log_dir, 'final_metrics_log.csv'), index=False)
    
    agent.save_checkpoint('final')        
   
    logging.info("Training Completed.")

# Test Function
def test(env, agent, num_targets, max_timesteps_per_target=300, test_targets=None):
    rewards_log = []
    metrics_log = []
    
    # Load test targets from saved state if not provided
    if not test_targets:
        saved_state = load_training_state()
        if saved_state and 'test_targets' in saved_state:
            test_targets = saved_state['test_targets']
        else:
            test_targets = CircuitEnvironment.generate_random_targets(num_targets=num_targets, seed=1)

    for i, target in enumerate(test_targets):
        node_features_tensor, edge_index, performance_metrics = env.reset(target)
        state = (node_features_tensor, edge_index)
    
        episode_rewards = []
        states = []
        actions = []
        values = []
        masks = []
        perf_metrics_log = []
        
        done = False
        t = 0
        while not done and t < max_timesteps_per_target:
            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)
            state = (next_node_features_tensor, next_edge_index)
    
            episode_rewards.append(reward)
            perf_metrics_log.append(perf_metrics)
                
            performance_metrics = previous_metrics
            t += 1
            
        total_reward = sum(episode_rewards)
        
        min_episode_reward = min(episode_rewards)
        max_episode_reward = max(episode_rewards)
        
        mean_episode_reward = total_reward / len(episode_rewards)
        deployment_accuracy = calculate_accuracy(perf_metrics_log)
        
        # Metrics calculations
        metrics_agg = {}
        for metric in ['gain', 'bw', 'ugf', 'pm']:
            metric_values = [m[metric] for m in perf_metrics_log]
            metrics_agg[f'{metric}_min'] = min(metric_values)
            metrics_agg[f'{metric}_max'] = max(metric_values)
            metrics_agg[f'{metric}_mean'] = sum(metric_values) / len(metric_values)
        
        print(f'Test: Total Reward: {total_reward}, Mean Episode Reward: {mean_episode_reward}, Min Episode Reward: {min_episode_reward}, Max Episode Reward: {max_episode_reward}, Deployment Accuracy: {deployment_accuracy}')
             
        # Logging
        logging.info(f"Test: Mean Episode Reward: {mean_episode_reward}, Min Episode Reward: {min_episode_reward}, Max Episode Reward: {max_episode_reward}, Deployment Accuracy: {deployment_accuracy}")   
        logging.info(f"Test: Episode {episode + 1} - Performance Metrics (min, max, mean): {metrics_agg}")
        
        rewards_log.append({'target_index': i, 'total_reward': total_reward})
        metrics_log.append({'target_index': i, 'Mean Episode Reward': mean_episode_reward, 'Min Episode Reward': min_episode_reward, 'Max Episode Reward': max_episode_reward, 'Deployment Accuracy': deployment_accuracy, **metrics_agg})
    
    log_dir = 'OPAMP_GNNRL_LTS_TrainTest_test3_log_file'
    os.makedirs(log_dir, exist_ok=True)
    pd.DataFrame(rewards_log).to_csv(os.path.join(log_dir, 'test_rewards_log.csv'), index=False)
    pd.DataFrame(metrics_log).to_csv(os.path.join(log_dir, 'test_metrics_log.csv'), index=False)
    
    logging.info("Testing Completed.")

if __name__ == "__main__":
    # Initialize environment parameters
    # Define the values for bounds, and target metrics
    bounds_low = np.array([0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 15e-6, 0.1e-12, 0.8])
    bounds_high = np.array([50e-6, 50e-6, 50e-6, 50e-6, 50e-6, 30e-6, 10e-12, 1.4])
    # Initialize Hyperparameters
    state_dim = 23
    action_space = 8
    lr_actor = 1e-4
    lr_critic = 3e-4
    gamma = 0.99
    gae_lambda = 0.95
    epsilon = 0.2
    policy_clip = 0.2
    epochs = 10
    entropy_coef = 0.01

    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)

    num_episodes = 1000
    max_timesteps = 300

    train(env, agent, num_episodes, max_timesteps)

    trained_state = load_training_state()
    if trained_state:
        agent.actor.load_state_dict(trained_state['actor_state_dict'])
        agent.critic.load_state_dict(trained_state['critic_state_dict'])

    num_test_targets = 20
    test(env, agent, num_test_targets, max_timesteps_per_target=300)