File size: 4,686 Bytes
9f8cf99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """Experience replay buffer for RL agents."""
from __future__ import annotations
from collections import deque
from dataclasses import dataclass
from typing import Optional, Tuple
import numpy as np
import torch
@dataclass
class Experience:
"""Single experience tuple for replay buffer."""
state: torch.Tensor
action: int
reward: float
next_state: torch.Tensor
done: bool
class ReplayBuffer:
"""Experience replay buffer for off-policy RL algorithms.
Stores and samples experience tuples (s, a, r, s', done) for training.
Supports prioritized experience replay with TD-error-based priorities.
"""
def __init__(self, capacity: int = 10000, prioritized: bool = True, alpha: float = 0.6):
"""Initialize replay buffer.
Args:
capacity: Maximum number of experiences to store
prioritized: Whether to use prioritized experience replay
alpha: Priority exponent (0 = uniform, 1 = full prioritization)
"""
self.capacity = capacity
self.prioritized = prioritized
self.alpha = alpha
self.buffer: deque[Experience] = deque(maxlen=capacity)
self.priorities: deque[float] = deque(maxlen=capacity)
self.position = 0
def push(self, experience: Experience, priority: Optional[float] = None) -> None:
"""Add experience to buffer.
Args:
experience: Experience tuple to store
priority: Priority for prioritized replay (default: max priority)
"""
self.buffer.append(experience)
if self.prioritized:
if priority is None:
# Use max priority for new experiences
priority = max(self.priorities) if self.priorities else 1.0
self.priorities.append(priority)
def sample(self, batch_size: int, beta: float = 0.4) -> Tuple[torch.Tensor, ...]:
"""Sample a batch of experiences.
Args:
batch_size: Number of experiences to sample
beta: Importance sampling exponent (0 = no correction, 1 = full correction)
Returns:
Tuple of (states, actions, rewards, next_states, dones, weights, indices)
"""
if len(self.buffer) < batch_size:
raise ValueError(f"Buffer contains {len(self.buffer)} experiences, need {batch_size}")
if self.prioritized:
# Sample based on priorities
priorities = np.array(self.priorities)
probs = priorities ** self.alpha
probs /= probs.sum()
indices = np.random.choice(len(self.buffer), batch_size, p=probs, replace=False)
# Compute importance sampling weights
weights = (len(self.buffer) * probs[indices]) ** (-beta)
weights /= weights.max() # Normalize
weights = torch.FloatTensor(weights)
else:
indices = np.random.choice(len(self.buffer), batch_size, replace=False)
weights = torch.ones(batch_size)
experiences = [self.buffer[idx] for idx in indices]
states = torch.stack([e.state for e in experiences])
# Actions might be discrete (Long) or continuous (Float). If they are tensors, stack them.
# If they are not tensors, convert to tensors.
if isinstance(experiences[0].action, torch.Tensor):
actions = torch.stack([e.action for e in experiences])
else:
try:
actions = torch.FloatTensor([e.action for e in experiences])
except (TypeError, ValueError):
actions = torch.stack([torch.tensor(e.action) for e in experiences])
rewards = torch.FloatTensor([e.reward for e in experiences])
next_states = torch.stack([e.next_state for e in experiences])
dones = torch.FloatTensor([e.done for e in experiences])
return states, actions, rewards, next_states, dones, weights, indices
def update_priorities(self, indices: np.ndarray, priorities: np.ndarray) -> None:
"""Update priorities for sampled experiences.
Args:
indices: Indices of sampled experiences
priorities: New priority values (typically TD errors)
"""
if not self.prioritized:
return
for idx, priority in zip(indices, priorities):
if 0 <= idx < len(self.priorities):
self.priorities[idx] = priority
def __len__(self) -> int:
return len(self.buffer)
def is_ready(self, batch_size: int) -> bool:
"""Check if buffer has enough experiences for sampling."""
return len(self.buffer) >= batch_size
|