yesterday my code was trained upto 13 episodes, while i resumed the same code from its saved state today, it again starts from the 11th episode, why the code doesnot loading the saved state from where it left(last completed training episode) on past iteration.


# 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)

    start_episode = 0    
    rewards_log = []
    losses_log = []
    metrics_log = []
    episode_lengths = []
    
    saved_state = load_training_state()
    if saved_state:
        start_episode = saved_state['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)
        
                # 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}")   
        
        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})
    
        episode_lengths.append(episode_length)
    
        agent.save_checkpoint(episode)
        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.")

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)