File size: 1,692 Bytes
a1fc554 | 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 | import sys
import os
ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
sys.path.append(ROOT_DIR)
os.chdir(ROOT_DIR)
import zarr
from diffusion_policy.common.replay_buffer import ReplayBuffer
def test():
import numpy as np
buff = ReplayBuffer.create_empty_numpy()
buff.add_episode({
'obs': np.zeros((100,10), dtype=np.float16)
})
buff.add_episode({
'obs': np.ones((50,10)),
'action': np.ones((50,2))
})
# buff.rechunk(256)
obs = buff.get_episode(0)
import numpy as np
buff = ReplayBuffer.create_empty_zarr()
buff.add_episode({
'obs': np.zeros((100,10), dtype=np.float16)
})
buff.add_episode({
'obs': np.ones((50,10)),
'action': np.ones((50,2))
})
obs = buff.get_episode(0)
buff.set_chunks({
'obs': (100,10),
'action': (100,2)
})
def test_real():
import os
dist_group = zarr.open(
os.path.expanduser('~/dev/diffusion_policy/data/pusht/pusht_cchi_v2.zarr'), 'r')
buff = ReplayBuffer.create_empty_numpy()
key, group = next(iter(dist_group.items()))
for key, group in dist_group.items():
buff.add_episode(group)
# out_path = os.path.expanduser('~/dev/diffusion_policy/data/pusht_cchi2_v2_replay.zarr')
out_path = os.path.expanduser('~/dev/diffusion_policy/data/test.zarr')
out_store = zarr.DirectoryStore(out_path)
buff.save_to_store(out_store)
buff = ReplayBuffer.copy_from_path(out_path, store=zarr.MemoryStore())
buff.pop_episode()
def test_pop():
buff = ReplayBuffer.create_from_path(
'/home/chengchi/dev/diffusion_policy/data/pusht_cchi_v3_replay.zarr',
mode='rw')
|