Spaces:
Sleeping
Sleeping
File size: 2,381 Bytes
208fbf8 | 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 | import numpy as np
import torch
import random
from typing import Dict, List, Tuple
class ReplayBuffer:
"""
A simple sequence replay buffer for Dreamer training.
Expects episodes of shape (seq_len, dim) to be added.
"""
def __init__(self, capacity: int = 10000, seq_len: int = 50):
self.capacity = capacity
self.seq_len = seq_len
self.episodes = []
def add_episode(self, obs: np.ndarray, actions: np.ndarray, rewards: np.ndarray):
"""
obs: (time, obs_dim)
actions: (time, action_dim)
rewards: (time,)
"""
if len(self.episodes) >= self.capacity:
self.episodes.pop(0)
self.episodes.append({
'obs': obs,
'actions': actions,
'rewards': rewards
})
def sample_batch(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Samples a batch of sequence chunks.
Returns tensors of shape (batch, seq_len, dim)
"""
if len(self.episodes) == 0:
raise ValueError("Replay buffer is empty.")
obs_batch, action_batch, reward_batch = [], [], []
for _ in range(batch_size):
ep_idx = random.randint(0, len(self.episodes) - 1)
ep = self.episodes[ep_idx]
ep_len = len(ep['obs'])
if ep_len <= self.seq_len:
start = 0
pad_len = self.seq_len - ep_len
obs = np.pad(ep['obs'], ((0, pad_len), (0, 0)), mode='edge')
acts = np.pad(ep['actions'], ((0, pad_len), (0, 0)), mode='edge')
rews = np.pad(ep['rewards'], (0, pad_len), mode='constant', constant_values=0)
else:
start = random.randint(0, ep_len - self.seq_len)
obs = ep['obs'][start:start + self.seq_len]
acts = ep['actions'][start:start + self.seq_len]
rews = ep['rewards'][start:start + self.seq_len]
obs_batch.append(obs)
action_batch.append(acts)
reward_batch.append(rews)
return (
torch.FloatTensor(np.stack(obs_batch)),
torch.FloatTensor(np.stack(action_batch)),
torch.FloatTensor(np.stack(reward_batch))
)
|