utils.py

# ============================================
# Utility functions and replay buffer
# ============================================

import random
import collections
import numpy as np
import torch


class ReplayBuffer:
    """Experience replay buffer for off-policy reinforcement learning."""

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

    def add(self, state, action, reward, next_state, done):
        """Add a transition to the buffer."""
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        """Sample a batch of transitions from memory."""
        transitions = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = zip(*transitions)

        return np.array(state), action, reward, np.array(next_state), done

    def size(self):
        """Return the number of stored transitions."""
        return len(self.buffer)


def set_seed(seed):
    """Set Python, NumPy, and PyTorch random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def moving_average(data, window_size=9):
    """Compute a simple moving average for smoothing data."""
    data = np.array(data, dtype=np.float32)

    if len(data) < window_size:
        return data

    weights = np.ones(window_size) / window_size
    return np.convolve(data, weights, mode="valid")