""" PALADIM Experience Replay Buffer ================================= Stores and samples past experiences to prevent catastrophic forgetting. The replay buffer maintains a diverse set of examples from previous tasks and mixes them with current training data during the rapid learning phase. """ import torch import torch.nn as nn from typing import Dict, List, Optional, Tuple, Any from collections import deque import random import numpy as np class ReplayBuffer: """ Experience Replay Buffer for Continual Learning. Stores (input, target) pairs from previous tasks and provides methods for sampling and mixing with current training data. Supports: - Reservoir sampling for uniform task distribution - Priority-based sampling based on loss/difficulty - Task-balanced sampling """ def __init__( self, capacity: int = 10000, sampling_strategy: str = "uniform", # "uniform", "priority", "balanced" ): """ Initialize replay buffer. Args: capacity: Maximum number of samples to store sampling_strategy: How to sample from buffer """ self.capacity = capacity self.sampling_strategy = sampling_strategy # Main storage self.buffer: List[Dict[str, Any]] = [] # Task-specific indices for balanced sampling self.task_indices: Dict[str, List[int]] = {} # Priority scores for priority sampling self.priorities: List[float] = [] # Reservoir sampling counter self.total_seen = 0 def add( self, samples: Dict[str, torch.Tensor], task_id: Optional[str] = None, priorities: Optional[List[float]] = None, ): """ Add samples to the buffer. Uses reservoir sampling when buffer is full to maintain uniform distribution over all seen samples. Args: samples: Dictionary with 'input_ids', 'attention_mask', 'labels', etc. task_id: Optional task identifier for balanced sampling priorities: Optional priority scores for each sample """ batch_size = samples['input_ids'].shape[0] for i in range(batch_size): sample = { k: v[i].clone().detach().cpu() for k, v in samples.items() if isinstance(v, torch.Tensor) } sample['_task_id'] = task_id priority = priorities[i] if priorities else 1.0 if len(self.buffer) < self.capacity: # Buffer not full - just add idx = len(self.buffer) self.buffer.append(sample) self.priorities.append(priority) # Track task indices if task_id: if task_id not in self.task_indices: self.task_indices[task_id] = [] self.task_indices[task_id].append(idx) else: # Reservoir sampling self.total_seen += 1 replace_idx = random.randint(0, self.total_seen) if replace_idx < self.capacity: # Remove old sample from task indices old_task = self.buffer[replace_idx].get('_task_id') if old_task and old_task in self.task_indices: if replace_idx in self.task_indices[old_task]: self.task_indices[old_task].remove(replace_idx) # Replace self.buffer[replace_idx] = sample self.priorities[replace_idx] = priority # Add new task index if task_id: if task_id not in self.task_indices: self.task_indices[task_id] = [] self.task_indices[task_id].append(replace_idx) def sample( self, batch_size: int, device: Optional[torch.device] = None, ) -> Optional[Dict[str, torch.Tensor]]: """ Sample a batch from the buffer. Args: batch_size: Number of samples to return device: Device to move tensors to Returns: Dictionary of batched tensors, or None if buffer is empty """ if len(self.buffer) == 0: return None batch_size = min(batch_size, len(self.buffer)) if self.sampling_strategy == "uniform": indices = random.sample(range(len(self.buffer)), batch_size) elif self.sampling_strategy == "priority": indices = self._priority_sample(batch_size) elif self.sampling_strategy == "balanced": indices = self._balanced_sample(batch_size) else: indices = random.sample(range(len(self.buffer)), batch_size) return self._collate(indices, device) def _priority_sample(self, batch_size: int) -> List[int]: """Sample based on priority scores.""" total_priority = sum(self.priorities) if total_priority == 0: return random.sample(range(len(self.buffer)), batch_size) probs = [p / total_priority for p in self.priorities] return np.random.choice( len(self.buffer), size=batch_size, replace=False, p=probs, ).tolist() def _balanced_sample(self, batch_size: int) -> List[int]: """Sample uniformly across tasks.""" if not self.task_indices: return random.sample(range(len(self.buffer)), batch_size) indices = [] tasks = list(self.task_indices.keys()) samples_per_task = max(1, batch_size // len(tasks)) for task_id in tasks: task_idx = self.task_indices[task_id] if task_idx: n = min(samples_per_task, len(task_idx)) indices.extend(random.sample(task_idx, n)) # Fill remaining with random samples remaining = batch_size - len(indices) if remaining > 0 and len(self.buffer) > len(indices): available = [i for i in range(len(self.buffer)) if i not in indices] indices.extend(random.sample(available, min(remaining, len(available)))) return indices[:batch_size] def _collate( self, indices: List[int], device: Optional[torch.device] = None, ) -> Dict[str, torch.Tensor]: """Collate samples into a batch.""" batch = {} samples = [self.buffer[i] for i in indices] # Get all keys except metadata keys = [k for k in samples[0].keys() if not k.startswith('_')] for key in keys: tensors = [s[key] for s in samples] batch[key] = torch.stack(tensors) if device: batch[key] = batch[key].to(device) return batch def update_priorities( self, indices: List[int], new_priorities: List[float], ): """Update priority scores for samples.""" for idx, priority in zip(indices, new_priorities): if 0 <= idx < len(self.priorities): self.priorities[idx] = priority def get_task_distribution(self) -> Dict[str, int]: """Get number of samples per task.""" return { task_id: len(indices) for task_id, indices in self.task_indices.items() } def __len__(self) -> int: return len(self.buffer) def clear(self): """Clear the buffer.""" self.buffer.clear() self.priorities.clear() self.task_indices.clear() self.total_seen = 0 class MixedDataLoader: """ DataLoader that mixes current task data with replay buffer samples. Yields batches that combine: - (1 - replay_ratio) from current dataloader - replay_ratio from replay buffer """ def __init__( self, current_dataloader, replay_buffer: ReplayBuffer, replay_ratio: float = 0.25, device: Optional[torch.device] = None, ): """ Initialize mixed dataloader. Args: current_dataloader: DataLoader for current task replay_buffer: ReplayBuffer with past experiences replay_ratio: Fraction of batch from replay (0.0-1.0) device: Device for tensors """ self.current_dataloader = current_dataloader self.replay_buffer = replay_buffer self.replay_ratio = replay_ratio self.device = device def __iter__(self): for batch in self.current_dataloader: # Move current batch to device if self.device: batch = { k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items() } # Get batch size batch_size = batch['input_ids'].shape[0] replay_size = int(batch_size * self.replay_ratio) if replay_size > 0 and len(self.replay_buffer) > 0: # Sample from replay buffer replay_batch = self.replay_buffer.sample(replay_size, self.device) if replay_batch: # Mix batches mixed_batch = {} for key in batch: if key in replay_batch: # Truncate current batch and append replay current = batch[key][:batch_size - replay_size] mixed_batch[key] = torch.cat([current, replay_batch[key]], dim=0) else: mixed_batch[key] = batch[key] yield mixed_batch continue yield batch def __len__(self): return len(self.current_dataloader) def create_replay_dataloader( current_dataloader, replay_buffer: ReplayBuffer, replay_ratio: float = 0.25, device: Optional[torch.device] = None, ) -> MixedDataLoader: """ Factory function to create a mixed dataloader. Args: current_dataloader: DataLoader for current task replay_buffer: ReplayBuffer instance replay_ratio: Fraction of replay samples (0.25 = 25%) device: Target device Returns: MixedDataLoader instance """ return MixedDataLoader( current_dataloader=current_dataloader, replay_buffer=replay_buffer, replay_ratio=replay_ratio, device=device, )