utils.py

import torch
from collections import deque
import random
import numpy as np

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def store(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        samples = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*samples)
        return (
            np.stack(states),
            np.array(actions),
            np.array(rewards),
            np.stack(next_states),
            np.array(dones),
        )

    def __len__(self):
        return len(self.buffer)

def save_model(agent, filepath):
    torch.save(agent.state_dict(), filepath)

def load_model(agent, filepath):
    agent.load_state_dict(torch.load(filepath))
    agent.eval()