I need you to ensure my below code is properly integrates my requirement of Data Storage & Logging in the existing code, if it is not properly integrated please update the necessary code for the proper synxhronized integration

1. Setting up a Logger:

import logging

# Setup basic configuration for logging
logging.basicConfig(filename=‘agent_training.log’, level=logging.INFO, format=‘%(asctime)s:%(levelname)s:%(message)s’)

2. Modifying the train Function to Log and Store Data:

import csv
import os

# initialize CSV files and writers at the beginning of your training session:
# Ensure the output directory exists
os.makedirs(‘training_output’, exist_ok=True)

# Open CSV files for writing
rewards_file = open(‘training_output/rewards.csv’, ‘w’, newline=‘’)
losses_file = open(‘training_output/losses.csv’, ‘w’, newline=‘’)

# Create CSV writers
rewards_writer = csv.writer(rewards_file)
losses_writer = csv.writer(losses_file)

# Write headers
rewards_writer.writerow([‘Episode’, ‘Immediate Reward’, ‘Cumulative Reward’])
losses_writer.writerow([‘Episode’, ‘Actor Loss’, ‘Critic Loss’])

# the train function to log and store data:
# Inside your episode loop after action and reward are received
rewards_writer.writerow([episode + 1, reward, sum(episode_rewards)])

# After computing losses in update_policy
losses_writer.writerow([episode + 1, actor_loss.item(), critic_loss.item()])

# capture actor_loss and critic_loss correctly in PPOAgent’s update_policy method
# Add logging at critical points
logging.info(f’Episode {episode + 1}, Total Reward: {total_reward}, Actor Loss: {actor_loss.item()}, Critic Loss: {critic_loss.item()}')


# close CSV files after the training loop.
rewards_file.close()
losses_file.close()


My code for optimization the analog circuit design parameters to get the desired performance metrics using reinforcement learning and graph neural network with PPO agent.
import pandas as pd
import logging
import os

# Basic configuration for logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

logger = logging.getLogger(__name__)

class PPOAgent:
    def __init__(self, gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef):
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.epsilon = epsilon
        self.policy_clip = policy_clip
        self.epochs = epochs
        self.entropy_coef = entropy_coef  # Dynamic entropy coefficient
        
        self.actor = Actor(gnn_model)
        self.critic = 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  # Assume continuous
        
        self.log_dir = 'OPAMP_GNNRL_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')))

    def select_action(self, state, performance_metrics):               
        
        return scaled_actions.detach().numpy(), log_probs.detach(), performance_metrics
    
    def compute_gae(self, next_value, rewards, masks, values):

        return returns

    def update_policy(self, prev_states, prev_actions, prev_log_probs, returns, advantages):
    
        return average_loss_actor, average_loss_critic


def train(env, agent, num_episodes, max_timesteps, batch_size):
    rewards_log = []
    losses_log = []
    episode_lengths = []
    
    for episode in range(num_episodes):        
        node_features_tensor, edge_feature_tensor, edge_index, performance_metrics = env.reset()
        state = (node_features_tensor, edge_feature_tensor, edge_index)

        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        masks = []

        for t in range(max_timesteps):
            # Selecting action based on the current state
            action, log_prob, perf_metrics = agent.select_action(state, performance_metrics)
            next_node_features_tensor, next_edge_feature_tensor, next_edge_index, reward, done, previous_metrics = env.step(action, perf_metrics)
            next_state = (next_node_features_tensor, next_edge_feature_tensor, next_edge_index)

            episode_rewards.append(reward)
            states.append(state)
            actions.append(action)
            log_probs.append(log_prob.item())
            values.append(agent.critic(state[0]).item())  # Ensure Critic can handle state format
            masks.append(1.0 - float(done))
            
            performance_metrics = previous_metrics
            state = next_state

            if done:
                break

        next_value = agent.critic(state[0]).item() if not done else 0        
        returns = agent.compute_gae(next_value, episode_rewards, masks, values)

        # Correcting format for updating the policy
        advantages = (torch.tensor(returns) - torch.tensor(values)).detach().numpy()
        returns = torch.tensor(returns).detach().numpy()
        # Update the policy with collected data
        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)
        
        # Log episode information
        total_reward = sum(episode_rewards)
        print(f'Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}')
        
        logger.info(f'Episode {episode} - Reward: {total_reward}, Length: {episode_length}')
        rewards_log.append({'episode': episode, 'reward': total_reward, 'length': episode_length})
        losses_log.append({'episode': episode, 'actor_loss': average_loss_actor, 'critic_loss': average_loss_critic})
        
        agent.save_checkpoint(episode)
        
        episode_lengths.append(episode_length)
        
        # Optional: Save episode logs to CSV every n episodes
        if episode % 5 == 0:
            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)
            
    # Final save for logs
    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)

    # Consider also saving final model checkpoints
    agent.save_checkpoint('final')

# Create the environment
env = CircuitEnvironment(server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content)

# Instantiate the GNN model
gnn_model = CustomGNN(in_channels, out_channels)

agent = PPOAgent(gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef)

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