File size: 2,688 Bytes
987ed1b | 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 | from typing import Optional, Dict
import numpy as np
from diffusion_policy.common.pref_replay_buffer import PrefReplayBuffer
import torch
def get_val_mask(n_episodes, val_ratio, seed=0):
val_mask = np.zeros(n_episodes, dtype=bool)
if val_ratio <= 0:
return val_mask
# have at least 1 episode for validation, and at least 1 episode for train
n_val = min(max(1, round(n_episodes * val_ratio)), n_episodes-1)
rng = np.random.default_rng(seed=seed)
val_idxs = rng.choice(n_episodes, size=n_val, replace=False)
val_mask[val_idxs] = True
return val_mask
class PrefSequenceSampler:
def __init__(self,
replay_buffer: PrefReplayBuffer,
sequence_length: int,
episode_mask: Optional[np.ndarray]=None,
keys: Optional[Dict[str, int]] = None,
):
"""
Initializes a sampler for the preference replay buffer.
Parameters:
- replay_buffer: PrefReplayBuffer instance from which to sample data.
- sequence_length: The length of sequences to sample.
- pad_before, pad_after: Padding before and after sequences (optional).
- keys: Optional dictionary to specify specific keys and limits on how much data to load.
- episode_mask: Mask indicating valid episodes for sampling.
"""
super().__init__()
assert sequence_length >= 1
if keys is None:
keys = list(replay_buffer.data.keys())
# Store generated indices
self.keys = keys
self.sequence_length = sequence_length
self.replay_buffer = replay_buffer
self.episode_mask = episode_mask
def __len__(self):
return np.sum(self.episode_mask)
def sample_sequence(self, idx: int) -> Dict[str, np.ndarray]:
"""
Samples the sequence of data based on the provided index (idx).
Parameters:
- idx: The index from which to sample an episode sequence.
Returns:
- A dictionary containing the sampled data for the specified keys and votes.
"""
indices = np.where(self.episode_mask)[0]
result = self.replay_buffer.get_pref_episode(indices[idx])
for key in result:
value = result[key]
if isinstance(value, np.ndarray):
result[key] = torch.from_numpy(value)
elif isinstance(value, (np.float32, np.float64, float, int)):
result[key] = torch.tensor(value, dtype=torch.float32)
else:
raise TypeError(f"Unsupported type {type(value)} for key '{key}'")
return result |