-
Notifications
You must be signed in to change notification settings - Fork 230
/
kitchen_lowdim_dataset.py
91 lines (80 loc) · 3.13 KB
/
kitchen_lowdim_dataset.py
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
from typing import Dict
import torch
import numpy as np
import copy
import pathlib
from diffusion_policy.common.pytorch_util import dict_apply
from diffusion_policy.common.replay_buffer import ReplayBuffer
from diffusion_policy.common.sampler import SequenceSampler, get_val_mask
from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer
from diffusion_policy.dataset.base_dataset import BaseLowdimDataset
class KitchenLowdimDataset(BaseLowdimDataset):
def __init__(self,
dataset_dir,
horizon=1,
pad_before=0,
pad_after=0,
seed=42,
val_ratio=0.0
):
super().__init__()
data_directory = pathlib.Path(dataset_dir)
observations = np.load(data_directory / "observations_seq.npy")
actions = np.load(data_directory / "actions_seq.npy")
masks = np.load(data_directory / "existence_mask.npy")
self.replay_buffer = ReplayBuffer.create_empty_numpy()
for i in range(len(masks)):
eps_len = int(masks[i].sum())
obs = observations[i,:eps_len].astype(np.float32)
action = actions[i,:eps_len].astype(np.float32)
data = {
'obs': obs,
'action': action
}
self.replay_buffer.add_episode(data)
val_mask = get_val_mask(
n_episodes=self.replay_buffer.n_episodes,
val_ratio=val_ratio,
seed=seed)
train_mask = ~val_mask
self.sampler = SequenceSampler(
replay_buffer=self.replay_buffer,
sequence_length=horizon,
pad_before=pad_before,
pad_after=pad_after,
episode_mask=train_mask)
self.train_mask = train_mask
self.horizon = horizon
self.pad_before = pad_before
self.pad_after = pad_after
def get_validation_dataset(self):
val_set = copy.copy(self)
val_set.sampler = SequenceSampler(
replay_buffer=self.replay_buffer,
sequence_length=self.horizon,
pad_before=self.pad_before,
pad_after=self.pad_after,
episode_mask=~self.train_mask
)
val_set.train_mask = ~self.train_mask
return val_set
def get_normalizer(self, mode='limits', **kwargs):
data = {
'obs': self.replay_buffer['obs'],
'action': self.replay_buffer['action']
}
if 'range_eps' not in kwargs:
# to prevent blowing up dims that barely change
kwargs['range_eps'] = 5e-2
normalizer = LinearNormalizer()
normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs)
return normalizer
def get_all_actions(self) -> torch.Tensor:
return torch.from_numpy(self.replay_buffer['action'])
def __len__(self) -> int:
return len(self.sampler)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
sample = self.sampler.sample_sequence(idx)
data = sample
torch_data = dict_apply(data, torch.from_numpy)
return torch_data