File size: 10,650 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 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 | from typing import Dict, List
import torch
import numpy as np
import zarr
import os
import shutil
from filelock import FileLock
from threadpoolctl import threadpool_limits
from omegaconf import OmegaConf
import cv2
import json
import hashlib
import copy
from diffusion_policy.common.pytorch_util import dict_apply
from diffusion_policy.dataset.base_dataset import BaseImageDataset
from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer
from diffusion_policy.common.replay_buffer import ReplayBuffer
from diffusion_policy.common.sampler import (
SequenceSampler, get_val_mask, downsample_mask)
from diffusion_policy.real_world.real_data_conversion import real_data_to_replay_buffer
from diffusion_policy.common.normalize_util import (
get_range_normalizer_from_stat,
get_image_range_normalizer,
get_identity_normalizer_from_stat,
array_to_stats
)
class RealPushTImageDataset(BaseImageDataset):
def __init__(self,
shape_meta: dict,
dataset_path: str,
horizon=1,
pad_before=0,
pad_after=0,
n_obs_steps=None,
n_latency_steps=0,
use_cache=False,
seed=42,
val_ratio=0.0,
max_train_episodes=None,
delta_action=False,
):
assert os.path.isdir(dataset_path)
replay_buffer = None
if use_cache:
# fingerprint shape_meta
shape_meta_json = json.dumps(OmegaConf.to_container(shape_meta), sort_keys=True)
shape_meta_hash = hashlib.md5(shape_meta_json.encode('utf-8')).hexdigest()
cache_zarr_path = os.path.join(dataset_path, shape_meta_hash + '.zarr.zip')
cache_lock_path = cache_zarr_path + '.lock'
print('Acquiring lock on cache.')
with FileLock(cache_lock_path):
if not os.path.exists(cache_zarr_path):
# cache does not exists
try:
print('Cache does not exist. Creating!')
replay_buffer = _get_replay_buffer(
dataset_path=dataset_path,
shape_meta=shape_meta,
store=zarr.MemoryStore()
)
print('Saving cache to disk.')
with zarr.ZipStore(cache_zarr_path) as zip_store:
replay_buffer.save_to_store(
store=zip_store
)
except Exception as e:
shutil.rmtree(cache_zarr_path)
raise e
else:
print('Loading cached ReplayBuffer from Disk.')
with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store:
replay_buffer = ReplayBuffer.copy_from_store(
src_store=zip_store, store=zarr.MemoryStore())
print('Loaded!')
else:
replay_buffer = _get_replay_buffer(
dataset_path=dataset_path,
shape_meta=shape_meta,
store=zarr.MemoryStore()
)
if delta_action:
# replace action as relative to previous frame
actions = replay_buffer['action'][:]
# support positions only at this time
assert actions.shape[1] <= 3
actions_diff = np.zeros_like(actions)
episode_ends = replay_buffer.episode_ends[:]
for i in range(len(episode_ends)):
start = 0
if i > 0:
start = episode_ends[i-1]
end = episode_ends[i]
# delta action is the difference between previous desired position and the current
# it should be scheduled at the previous timestep for the current timestep
# to ensure consistency with positional mode
actions_diff[start+1:end] = np.diff(actions[start:end], axis=0)
replay_buffer['action'][:] = actions_diff
rgb_keys = list()
lowdim_keys = list()
obs_shape_meta = shape_meta['obs']
for key, attr in obs_shape_meta.items():
type = attr.get('type', 'low_dim')
if type == 'rgb':
rgb_keys.append(key)
elif type == 'low_dim':
lowdim_keys.append(key)
key_first_k = dict()
if n_obs_steps is not None:
# only take first k obs from images
for key in rgb_keys + lowdim_keys:
key_first_k[key] = n_obs_steps
val_mask = get_val_mask(
n_episodes=replay_buffer.n_episodes,
val_ratio=val_ratio,
seed=seed)
train_mask = ~val_mask
train_mask = downsample_mask(
mask=train_mask,
max_n=max_train_episodes,
seed=seed)
sampler = SequenceSampler(
replay_buffer=replay_buffer,
sequence_length=horizon+n_latency_steps,
pad_before=pad_before,
pad_after=pad_after,
episode_mask=train_mask,
key_first_k=key_first_k)
self.replay_buffer = replay_buffer
self.sampler = sampler
self.shape_meta = shape_meta
self.rgb_keys = rgb_keys
self.lowdim_keys = lowdim_keys
self.n_obs_steps = n_obs_steps
self.val_mask = val_mask
self.horizon = horizon
self.n_latency_steps = n_latency_steps
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+self.n_latency_steps,
pad_before=self.pad_before,
pad_after=self.pad_after,
episode_mask=self.val_mask
)
val_set.val_mask = ~self.val_mask
return val_set
def get_normalizer(self, **kwargs) -> LinearNormalizer:
normalizer = LinearNormalizer()
# action
normalizer['action'] = SingleFieldLinearNormalizer.create_fit(
self.replay_buffer['action'])
# obs
for key in self.lowdim_keys:
normalizer[key] = SingleFieldLinearNormalizer.create_fit(
self.replay_buffer[key])
# image
for key in self.rgb_keys:
normalizer[key] = get_image_range_normalizer()
return normalizer
def get_all_actions(self) -> torch.Tensor:
return torch.from_numpy(self.replay_buffer['action'])
def __len__(self):
return len(self.sampler)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
threadpool_limits(1)
data = self.sampler.sample_sequence(idx)
# to save RAM, only return first n_obs_steps of OBS
# since the rest will be discarded anyway.
# when self.n_obs_steps is None
# this slice does nothing (takes all)
T_slice = slice(self.n_obs_steps)
obs_dict = dict()
for key in self.rgb_keys:
# move channel last to channel first
# T,H,W,C
# convert uint8 image to float32
obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1
).astype(np.float32) / 255.
# T,C,H,W
# save ram
del data[key]
for key in self.lowdim_keys:
obs_dict[key] = data[key][T_slice].astype(np.float32)
# save ram
del data[key]
action = data['action'].astype(np.float32)
# handle latency by dropping first n_latency_steps action
# observations are already taken care of by T_slice
if self.n_latency_steps > 0:
action = action[self.n_latency_steps:]
torch_data = {
'obs': dict_apply(obs_dict, torch.from_numpy),
'action': torch.from_numpy(action)
}
return torch_data
def zarr_resize_index_last_dim(zarr_arr, idxs):
actions = zarr_arr[:]
actions = actions[...,idxs]
zarr_arr.resize(zarr_arr.shape[:-1] + (len(idxs),))
zarr_arr[:] = actions
return zarr_arr
def _get_replay_buffer(dataset_path, shape_meta, store):
# parse shape meta
rgb_keys = list()
lowdim_keys = list()
out_resolutions = dict()
lowdim_shapes = dict()
obs_shape_meta = shape_meta['obs']
for key, attr in obs_shape_meta.items():
type = attr.get('type', 'low_dim')
shape = tuple(attr.get('shape'))
if type == 'rgb':
rgb_keys.append(key)
c,h,w = shape
out_resolutions[key] = (w,h)
elif type == 'low_dim':
lowdim_keys.append(key)
lowdim_shapes[key] = tuple(shape)
if 'pose' in key:
assert tuple(shape) in [(2,),(6,)]
action_shape = tuple(shape_meta['action']['shape'])
assert action_shape in [(2,),(6,)]
# load data
cv2.setNumThreads(1)
with threadpool_limits(1):
replay_buffer = real_data_to_replay_buffer(
dataset_path=dataset_path,
out_store=store,
out_resolutions=out_resolutions,
lowdim_keys=lowdim_keys + ['action'],
image_keys=rgb_keys
)
# transform lowdim dimensions
if action_shape == (2,):
# 2D action space, only controls X and Y
zarr_arr = replay_buffer['action']
zarr_resize_index_last_dim(zarr_arr, idxs=[0,1])
for key, shape in lowdim_shapes.items():
if 'pose' in key and shape == (2,):
# only take X and Y
zarr_arr = replay_buffer[key]
zarr_resize_index_last_dim(zarr_arr, idxs=[0,1])
return replay_buffer
def test():
import hydra
from omegaconf import OmegaConf
OmegaConf.register_new_resolver("eval", eval, replace=True)
with hydra.initialize('../diffusion_policy/config'):
cfg = hydra.compose('train_robomimic_real_image_workspace')
OmegaConf.resolve(cfg)
dataset = hydra.utils.instantiate(cfg.task.dataset)
from matplotlib import pyplot as plt
normalizer = dataset.get_normalizer()
nactions = normalizer['action'].normalize(dataset.replay_buffer['action'][:])
diff = np.diff(nactions, axis=0)
dists = np.linalg.norm(np.diff(nactions, axis=0), axis=-1)
_ = plt.hist(dists, bins=100); plt.title('real action velocity')
|