diff --git a/lerobot/common/datasets/aloha.py b/lerobot/common/datasets/aloha.py index 102de08e..639acf1f 100644 --- a/lerobot/common/datasets/aloha.py +++ b/lerobot/common/datasets/aloha.py @@ -158,6 +158,7 @@ class AlohaDataset(torch.utils.data.Dataset): self.data_ids_per_episode = {} ep_dicts = [] + idx0 = idx1 = 0 logging.info("Initialize and feed offline buffer") for ep_id in tqdm.tqdm(range(NUM_EPISODES[self.dataset_id])): ep_path = raw_dir / f"episode_{ep_id}.hdf5" @@ -165,7 +166,7 @@ class AlohaDataset(torch.utils.data.Dataset): num_frames = ep["/action"].shape[0] # last step of demonstration is considered done - done = torch.zeros(num_frames, 1, dtype=torch.bool) + done = torch.zeros(num_frames, dtype=torch.bool) done[-1] = True state = torch.from_numpy(ep["/observations/qpos"][:]) @@ -192,6 +193,14 @@ class AlohaDataset(torch.utils.data.Dataset): ep_dicts.append(ep_dict) + idx1 += num_frames + + assert isinstance(ep_id, int) + self.data_ids_per_episode[ep_id] = torch.arange(idx0, idx1, 1) + assert len(self.data_ids_per_episode[ep_id]) == num_frames + + idx0 = idx1 + self.data_dict = {} keys = ep_dicts[0].keys() diff --git a/lerobot/common/datasets/pusht.py b/lerobot/common/datasets/pusht.py index 9088fdf4..b468637e 100644 --- a/lerobot/common/datasets/pusht.py +++ b/lerobot/common/datasets/pusht.py @@ -193,8 +193,6 @@ class PushtDataset(torch.utils.data.Dataset): idx0 = 0 for episode_id in tqdm.tqdm(range(num_episodes)): idx1 = dataset_dict.meta["episode_ends"][episode_id] - # to create test artifact - # idx1 = 51 num_frames = idx1 - idx0 @@ -207,9 +205,9 @@ class PushtDataset(torch.utils.data.Dataset): block_pos = state[:, 2:4] block_angle = state[:, 4] - reward = torch.zeros(num_frames, 1) - success = torch.zeros(num_frames, 1, dtype=torch.bool) - done = torch.zeros(num_frames, 1, dtype=torch.bool) + reward = torch.zeros(num_frames) + success = torch.zeros(num_frames, dtype=torch.bool) + done = torch.zeros(num_frames, dtype=torch.bool) for i in range(num_frames): space = pymunk.Space() space.gravity = 0, 0 diff --git a/lerobot/common/datasets/utils.py b/lerobot/common/datasets/utils.py index 18b091cd..3b4aacfc 100644 --- a/lerobot/common/datasets/utils.py +++ b/lerobot/common/datasets/utils.py @@ -92,11 +92,11 @@ def load_data_with_delta_timestamps( # TODO(rcadene): synchronize timestamps + interpolation if needed - tol = 0.02 + tol = 0.04 is_pad = min_ > tol assert is_contiguously_true_or_false(is_pad), ( - "One or several timestamps unexpectedly violate the tolerance." + f"One or several timestamps unexpectedly violate the tolerance ({min_} > {tol=})." "This might be due to synchronization issues with timestamps during data collection." ) diff --git a/lerobot/common/policies/tdmpc/policy.py b/lerobot/common/policies/tdmpc/policy.py index f763dbc6..04aa5b11 100644 --- a/lerobot/common/policies/tdmpc/policy.py +++ b/lerobot/common/policies/tdmpc/policy.py @@ -429,7 +429,7 @@ class TDMPCPolicy(nn.Module): batch[key] = batch[key].transpose(1, 0) action = batch["action"] - reward = batch["next.reward"][:, :, None] # add extra channel dimension + reward = batch["next.reward"] # idxs = batch["index"] # TODO(rcadene): use idxs to update sampling weights done = torch.zeros_like(reward, dtype=torch.bool, device=reward.device) mask = torch.ones_like(reward, dtype=torch.bool, device=reward.device) diff --git a/lerobot/configs/policy/diffusion.yaml b/lerobot/configs/policy/diffusion.yaml index c3bebe2d..6da62e10 100644 --- a/lerobot/configs/policy/diffusion.yaml +++ b/lerobot/configs/policy/diffusion.yaml @@ -63,9 +63,9 @@ policy: grad_clip_norm: 10 delta_timestamps: - observation.image: [-.1, 0] - observation.state: [-.1, 0] - action: [-.1, 0, .1, .2, .3, .4, .5, .6, .7, .8, .9, 1.0, 1.1, 1.2, 1.3, 1.4] + observation.image: [-0.1, 0] + observation.state: [-0.1, 0] + action: [-0.1, 0, .1, .2, .3, .4, .5, .6, .7, .8, .9, 1.0, 1.1, 1.2, 1.3, 1.4] noise_scheduler: _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler diff --git a/tests/data/aloha_sim_insertion_human/data_dict.pth b/tests/data/aloha_sim_insertion_human/data_dict.pth new file mode 100644 index 00000000..1370c9ea Binary files /dev/null and b/tests/data/aloha_sim_insertion_human/data_dict.pth differ diff --git a/tests/data/aloha_sim_insertion_human/data_ids_per_episode.pth b/tests/data/aloha_sim_insertion_human/data_ids_per_episode.pth new file mode 100644 index 00000000..a1d481dd Binary files /dev/null and b/tests/data/aloha_sim_insertion_human/data_ids_per_episode.pth differ diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/action.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/action.memmap deleted file mode 100644 index f64b2989..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/action.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d789deddb081a9f4b626342391de8f48949d38fb5fdead87b5c0737b46c0877a -size 2800 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/episode.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/episode.memmap deleted file mode 100644 index af9fb07f..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/episode.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a12e561363385e9dfeeab326368731c030ed4b374e7f5897ac819159d2884c5 -size 400 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/frame_id.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/frame_id.memmap deleted file mode 100644 index dc2f585c..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/frame_id.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c202d9cfc7858fd49d522047e16948359bbbb2eda2d3825d552e45a78d5f8585 -size 400 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/meta.json b/tests/data/aloha_sim_insertion_human/replay_buffer/meta.json deleted file mode 100644 index 2a0cf0a2..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"action": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/next/done.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/next/done.memmap deleted file mode 100644 index 44fd709f..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/next/done.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51 -size 50 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/next/meta.json b/tests/data/aloha_sim_insertion_human/replay_buffer/next/meta.json deleted file mode 100644 index 3bfa9bd7..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/next/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"done": {"device": "cpu", "shape": [50, 1], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/image/meta.json b/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/image/top.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/image/top.memmap deleted file mode 100644 index d3d8bd1c..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5c632e3cb06be729e5d673e3ecca1d6f6527b0f48cfe3dc03d7eea4f9eb3bbd7 -size 46080000 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/meta.json b/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/state.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/state.memmap deleted file mode 100644 index 1f087a60..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/next/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e231f2e07e1cd030137ea2e938b570b112db2c694c6d21b37ceb8f8559e19088 -size 2800 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/image/meta.json b/tests/data/aloha_sim_insertion_human/replay_buffer/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/image/top.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/observation/image/top.memmap deleted file mode 100644 index 00c0b783..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a1ba64c89f4fcf9135fe34c26abf582dd5f0d573506db5c96af3ffe40a52c818 -size 46080000 diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/meta.json b/tests/data/aloha_sim_insertion_human/replay_buffer/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/state.memmap b/tests/data/aloha_sim_insertion_human/replay_buffer/observation/state.memmap deleted file mode 100644 index a1131179..00000000 --- a/tests/data/aloha_sim_insertion_human/replay_buffer/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:85405686bc065c6ab6c915907920a0391a57cf097b74de058a8c30be0548ade5 -size 2800 diff --git a/tests/data/aloha_sim_insertion_human/stats.pth b/tests/data/aloha_sim_insertion_human/stats.pth index 87b18a24..a7b9248f 100644 Binary files a/tests/data/aloha_sim_insertion_human/stats.pth and b/tests/data/aloha_sim_insertion_human/stats.pth differ diff --git a/tests/data/aloha_sim_insertion_scripted/data_dict.pth b/tests/data/aloha_sim_insertion_scripted/data_dict.pth new file mode 100644 index 00000000..00c9f335 Binary files /dev/null and b/tests/data/aloha_sim_insertion_scripted/data_dict.pth differ diff --git a/tests/data/aloha_sim_insertion_scripted/data_ids_per_episode.pth b/tests/data/aloha_sim_insertion_scripted/data_ids_per_episode.pth new file mode 100644 index 00000000..a1d481dd Binary files /dev/null and b/tests/data/aloha_sim_insertion_scripted/data_ids_per_episode.pth differ diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/action.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/action.memmap deleted file mode 100644 index e4068b75..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/action.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1f5fe053b760e8471885b82c10f4a6ea40874098036337ae5cc300c4775546be -size 2800 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/episode.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/episode.memmap deleted file mode 100644 index af9fb07f..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/episode.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a12e561363385e9dfeeab326368731c030ed4b374e7f5897ac819159d2884c5 -size 400 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/frame_id.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/frame_id.memmap deleted file mode 100644 index dc2f585c..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/frame_id.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c202d9cfc7858fd49d522047e16948359bbbb2eda2d3825d552e45a78d5f8585 -size 400 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/meta.json b/tests/data/aloha_sim_insertion_scripted/replay_buffer/meta.json deleted file mode 100644 index 2a0cf0a2..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"action": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/done.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/done.memmap deleted file mode 100644 index 44fd709f..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/done.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51 -size 50 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/meta.json b/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/meta.json deleted file mode 100644 index 3bfa9bd7..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"done": {"device": "cpu", "shape": [50, 1], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/image/meta.json b/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/image/top.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/image/top.memmap deleted file mode 100644 index 83911729..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:daed2bb10498ba2557983d0d7e89399882fea7585e7ceff910e23c621bfdbf88 -size 46080000 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/meta.json b/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/state.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/state.memmap deleted file mode 100644 index aef69da0..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/next/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:bbad0302af70112ee312efe0eb0f44a2f1c8f6c5ef82ea4fb34625cdafbef057 -size 2800 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/image/meta.json b/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/image/top.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/image/top.memmap deleted file mode 100644 index f9f0a759..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:aba55ebb9dd004bf68444b9ebf024ed7713436099c06a0b8e541100ecbc69290 -size 46080000 diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/meta.json b/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/state.memmap b/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/state.memmap deleted file mode 100644 index 91875055..00000000 --- a/tests/data/aloha_sim_insertion_scripted/replay_buffer/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:dd4e7e14abf57561ca9839c910581266be90956e41bfb3bb21362ea0c321e77d -size 2800 diff --git a/tests/data/aloha_sim_insertion_scripted/stats.pth b/tests/data/aloha_sim_insertion_scripted/stats.pth index 7d149ca4..990d4647 100644 Binary files a/tests/data/aloha_sim_insertion_scripted/stats.pth and b/tests/data/aloha_sim_insertion_scripted/stats.pth differ diff --git a/tests/data/aloha_sim_transfer_cube_human/data_dict.pth b/tests/data/aloha_sim_transfer_cube_human/data_dict.pth new file mode 100644 index 00000000..ab851779 Binary files /dev/null and b/tests/data/aloha_sim_transfer_cube_human/data_dict.pth differ diff --git a/tests/data/aloha_sim_transfer_cube_human/data_ids_per_episode.pth b/tests/data/aloha_sim_transfer_cube_human/data_ids_per_episode.pth new file mode 100644 index 00000000..a1d481dd Binary files /dev/null and b/tests/data/aloha_sim_transfer_cube_human/data_ids_per_episode.pth differ diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/action.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/action.memmap deleted file mode 100644 index 9b4fef33..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/action.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:14fed0eed3d529a8ac0dd25a6d41585020772d02f9137fc9d604713b2f0f7076 -size 2800 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/episode.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/episode.memmap deleted file mode 100644 index af9fb07f..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/episode.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a12e561363385e9dfeeab326368731c030ed4b374e7f5897ac819159d2884c5 -size 400 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/frame_id.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/frame_id.memmap deleted file mode 100644 index dc2f585c..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/frame_id.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c202d9cfc7858fd49d522047e16948359bbbb2eda2d3825d552e45a78d5f8585 -size 400 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/meta.json b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/meta.json deleted file mode 100644 index 2a0cf0a2..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"action": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/done.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/done.memmap deleted file mode 100644 index 44fd709f..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/done.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51 -size 50 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/meta.json b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/meta.json deleted file mode 100644 index 3bfa9bd7..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"done": {"device": "cpu", "shape": [50, 1], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/image/meta.json b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/image/top.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/image/top.memmap deleted file mode 100644 index cd2e7c06..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2f713ea7fc19e592ea409a5e0bdfde403e5b86f834cbabe3463b791e8437fafc -size 46080000 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/meta.json b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/state.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/state.memmap deleted file mode 100644 index 37feaad6..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/next/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8c103e2c9d63c9f7cf9645bd24d9a2c4e8e08825dc75e230ebc793b8f9c213b0 -size 2800 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/image/meta.json b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/image/top.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/image/top.memmap deleted file mode 100644 index 1188590c..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7dbf4aa01b184d0eaa21ea999078d7cff86e1ca484a109614176fdc49f1ee05c -size 46080000 diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/meta.json b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/state.memmap b/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/state.memmap deleted file mode 100644 index 9ef4cfd6..00000000 --- a/tests/data/aloha_sim_transfer_cube_human/replay_buffer/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4fa0b9c870d4615037b6fee9e9e85e54d84352e173f2c7c1035232272fe2a3dd -size 2800 diff --git a/tests/data/aloha_sim_transfer_cube_human/stats.pth b/tests/data/aloha_sim_transfer_cube_human/stats.pth index 22f3e4d9..1ae356e3 100644 Binary files a/tests/data/aloha_sim_transfer_cube_human/stats.pth and b/tests/data/aloha_sim_transfer_cube_human/stats.pth differ diff --git a/tests/data/aloha_sim_transfer_cube_scripted/data_dict.pth b/tests/data/aloha_sim_transfer_cube_scripted/data_dict.pth new file mode 100644 index 00000000..bd308bb0 Binary files /dev/null and b/tests/data/aloha_sim_transfer_cube_scripted/data_dict.pth differ diff --git a/tests/data/aloha_sim_transfer_cube_scripted/data_ids_per_episode.pth b/tests/data/aloha_sim_transfer_cube_scripted/data_ids_per_episode.pth new file mode 100644 index 00000000..a1d481dd Binary files /dev/null and b/tests/data/aloha_sim_transfer_cube_scripted/data_ids_per_episode.pth differ diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/action.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/action.memmap deleted file mode 100644 index 8ac0c726..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/action.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c0e199a82e2b7462e84406dbced5448a99f1dad9ce172771dfc3feb4b8597115 -size 2800 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/episode.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/episode.memmap deleted file mode 100644 index af9fb07f..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/episode.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a12e561363385e9dfeeab326368731c030ed4b374e7f5897ac819159d2884c5 -size 400 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/frame_id.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/frame_id.memmap deleted file mode 100644 index dc2f585c..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/frame_id.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c202d9cfc7858fd49d522047e16948359bbbb2eda2d3825d552e45a78d5f8585 -size 400 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/meta.json b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/meta.json deleted file mode 100644 index 2a0cf0a2..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"action": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/done.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/done.memmap deleted file mode 100644 index 44fd709f..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/done.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51 -size 50 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/meta.json b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/meta.json deleted file mode 100644 index 3bfa9bd7..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"done": {"device": "cpu", "shape": [50, 1], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/image/meta.json b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/image/top.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/image/top.memmap deleted file mode 100644 index 8e5f533e..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:30b44d38cc4d68e06c716a875d39cbdbeacbfdc1657d6366f58c279efd27c52b -size 46080000 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/meta.json b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/state.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/state.memmap deleted file mode 100644 index e88320d1..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/next/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9f484e7ea4f5f612dd53ee2c0f7891b8f7b2168a54fc81941ac2f2447260c294 -size 2800 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/image/meta.json b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/image/meta.json deleted file mode 100644 index cb29a5ab..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/image/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"top": {"device": "cpu", "shape": [50, 3, 480, 640], "dtype": "torch.uint8"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/image/top.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/image/top.memmap deleted file mode 100644 index d415da0a..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/image/top.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:09600206f56cc5b52dfb896204b0044c4e830da368da141d7bd10e52181f6835 -size 46080000 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/meta.json b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/meta.json deleted file mode 100644 index 65ce1ca2..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"state": {"device": "cpu", "shape": [50, 14], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/state.memmap b/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/state.memmap deleted file mode 100644 index be3436fb..00000000 --- a/tests/data/aloha_sim_transfer_cube_scripted/replay_buffer/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:60dcb547cf9a6372b78a455217a2408b6bece4371fba1df2a302b334d45c42a8 -size 2800 diff --git a/tests/data/aloha_sim_transfer_cube_scripted/stats.pth b/tests/data/aloha_sim_transfer_cube_scripted/stats.pth index 63465344..71547f09 100644 Binary files a/tests/data/aloha_sim_transfer_cube_scripted/stats.pth and b/tests/data/aloha_sim_transfer_cube_scripted/stats.pth differ diff --git a/tests/data/pusht/data_dict.pth b/tests/data/pusht/data_dict.pth new file mode 100644 index 00000000..40d96a51 Binary files /dev/null and b/tests/data/pusht/data_dict.pth differ diff --git a/tests/data/pusht/data_ids_per_episode.pth b/tests/data/pusht/data_ids_per_episode.pth new file mode 100644 index 00000000..a1d481dd Binary files /dev/null and b/tests/data/pusht/data_ids_per_episode.pth differ diff --git a/tests/data/pusht/replay_buffer/action.memmap b/tests/data/pusht/replay_buffer/action.memmap deleted file mode 100644 index f4127fb1..00000000 --- a/tests/data/pusht/replay_buffer/action.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ba17d8e5c30151ea5f7f6fc31f19e12a68ce2113774b74c8aca0c7ef962a75f4 -size 400 diff --git a/tests/data/pusht/replay_buffer/episode.memmap b/tests/data/pusht/replay_buffer/episode.memmap deleted file mode 100644 index af9fb07f..00000000 --- a/tests/data/pusht/replay_buffer/episode.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a12e561363385e9dfeeab326368731c030ed4b374e7f5897ac819159d2884c5 -size 400 diff --git a/tests/data/pusht/replay_buffer/frame_id.memmap b/tests/data/pusht/replay_buffer/frame_id.memmap deleted file mode 100644 index dc2f585c..00000000 --- a/tests/data/pusht/replay_buffer/frame_id.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c202d9cfc7858fd49d522047e16948359bbbb2eda2d3825d552e45a78d5f8585 -size 400 diff --git a/tests/data/pusht/replay_buffer/meta.json b/tests/data/pusht/replay_buffer/meta.json deleted file mode 100644 index 6f7c4218..00000000 --- a/tests/data/pusht/replay_buffer/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"action": {"device": "cpu", "shape": [50, 2], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/pusht/replay_buffer/next/done.memmap b/tests/data/pusht/replay_buffer/next/done.memmap deleted file mode 100644 index 44fd709f..00000000 --- a/tests/data/pusht/replay_buffer/next/done.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51 -size 50 diff --git a/tests/data/pusht/replay_buffer/next/meta.json b/tests/data/pusht/replay_buffer/next/meta.json deleted file mode 100644 index b29a9ff7..00000000 --- a/tests/data/pusht/replay_buffer/next/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"reward": {"device": "cpu", "shape": [50, 1], "dtype": "torch.float32"}, "done": {"device": "cpu", "shape": [50, 1], "dtype": "torch.bool"}, "success": {"device": "cpu", "shape": [50, 1], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/pusht/replay_buffer/next/observation/image.memmap b/tests/data/pusht/replay_buffer/next/observation/image.memmap deleted file mode 100644 index 68634378..00000000 --- a/tests/data/pusht/replay_buffer/next/observation/image.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ff6a3748c8223a82e54c61442df7b8baf478a20497ee2353645a1e9ccd765162 -size 5529600 diff --git a/tests/data/pusht/replay_buffer/next/observation/meta.json b/tests/data/pusht/replay_buffer/next/observation/meta.json deleted file mode 100644 index 57e0edea..00000000 --- a/tests/data/pusht/replay_buffer/next/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"image": {"device": "cpu", "shape": [50, 3, 96, 96], "dtype": "torch.float32"}, "state": {"device": "cpu", "shape": [50, 2], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/pusht/replay_buffer/next/observation/state.memmap b/tests/data/pusht/replay_buffer/next/observation/state.memmap deleted file mode 100644 index 8dd28f2a..00000000 --- a/tests/data/pusht/replay_buffer/next/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fad4ece6d5fd66bbafa34f6ff383c483410082b8d7d4f4616808c3c458ce1d43 -size 400 diff --git a/tests/data/pusht/replay_buffer/next/reward.memmap b/tests/data/pusht/replay_buffer/next/reward.memmap deleted file mode 100644 index 109ed5ad..00000000 --- a/tests/data/pusht/replay_buffer/next/reward.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6d9c54dee5660c46886f32d80e57e9dd0ffa57ee0cd2a762b036d9c8e0c3a33a -size 200 diff --git a/tests/data/pusht/replay_buffer/next/success.memmap b/tests/data/pusht/replay_buffer/next/success.memmap deleted file mode 100644 index 44fd709f..00000000 --- a/tests/data/pusht/replay_buffer/next/success.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc2786e1f9910a9d811400edcddaf7075195f7a16b216dcbefba3bc7c4f2ae51 -size 50 diff --git a/tests/data/pusht/replay_buffer/observation/image.memmap b/tests/data/pusht/replay_buffer/observation/image.memmap deleted file mode 100644 index 42c86ef0..00000000 --- a/tests/data/pusht/replay_buffer/observation/image.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4bbde5cfd8cff9fd9fc6c9a57177f6fd31c8a03cf853b7d2234312f38380b0ba -size 5529600 diff --git a/tests/data/pusht/replay_buffer/observation/meta.json b/tests/data/pusht/replay_buffer/observation/meta.json deleted file mode 100644 index 57e0edea..00000000 --- a/tests/data/pusht/replay_buffer/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"image": {"device": "cpu", "shape": [50, 3, 96, 96], "dtype": "torch.float32"}, "state": {"device": "cpu", "shape": [50, 2], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/pusht/replay_buffer/observation/state.memmap b/tests/data/pusht/replay_buffer/observation/state.memmap deleted file mode 100644 index 3ac8e4ab..00000000 --- a/tests/data/pusht/replay_buffer/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:67c7e39090a16546fb1eade833d704f26464d574d7e431415f828159a154d2bf -size 400 diff --git a/tests/data/pusht/stats.pth b/tests/data/pusht/stats.pth index d7107185..636985fd 100644 Binary files a/tests/data/pusht/stats.pth and b/tests/data/pusht/stats.pth differ diff --git a/tests/data/xarm_lift_medium/data_dict.pth b/tests/data/xarm_lift_medium/data_dict.pth new file mode 100644 index 00000000..5c166576 Binary files /dev/null and b/tests/data/xarm_lift_medium/data_dict.pth differ diff --git a/tests/data/xarm_lift_medium/data_ids_per_episode.pth b/tests/data/xarm_lift_medium/data_ids_per_episode.pth new file mode 100644 index 00000000..21095017 Binary files /dev/null and b/tests/data/xarm_lift_medium/data_ids_per_episode.pth differ diff --git a/tests/data/xarm_lift_medium/replay_buffer/action.memmap b/tests/data/xarm_lift_medium/replay_buffer/action.memmap deleted file mode 100644 index c90afbe9..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/action.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:10ec2f944de18f1a2aa3fc2f9a4185c03e3a5afc31148c85c98b58602ac4186e -size 800 diff --git a/tests/data/xarm_lift_medium/replay_buffer/episode.memmap b/tests/data/xarm_lift_medium/replay_buffer/episode.memmap deleted file mode 100644 index 7924f028..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/episode.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1a589cba6bf0dfce138110864b6509508a804d7ea5c519d0a3cd67c4a87fa2d0 -size 200 diff --git a/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap b/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap deleted file mode 100644 index a633d346..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6afe7098f30bdc8564526517c085e62613f6cb67194153840567974cfa6f3815 -size 400 diff --git a/tests/data/xarm_lift_medium/replay_buffer/meta.json b/tests/data/xarm_lift_medium/replay_buffer/meta.json deleted file mode 100644 index 33dc932c..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"action": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int32"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap deleted file mode 100644 index cf5e9cca..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:dab3a9712c413c4bfcd91c645752ab981306b23d25bcd4da4c422911574673f7 -size 50 diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/meta.json b/tests/data/xarm_lift_medium/replay_buffer/next/meta.json deleted file mode 100644 index d69cadad..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/next/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"reward": {"device": "cpu", "shape": [50], "dtype": "torch.float32"}, "done": {"device": "cpu", "shape": [50], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap deleted file mode 100644 index 462d0117..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d6f9d1422ce3764e7253f70ed4da278f0c07fafef0d5386479f09d6b6b9d8259 -size 1058400 diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json b/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json deleted file mode 100644 index b13b8ec9..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"image": {"device": "cpu", "shape": [50, 3, 84, 84], "dtype": "torch.uint8"}, "state": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap deleted file mode 100644 index 1dbe6024..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:52e7c1a3c4fb2423b195e66ffee2c9e23b3ea0ad7c8bfc4dec30a35c65cadcbb -size 800 diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap deleted file mode 100644 index 9ff5d5a1..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c4dbe8ea1966e5cc6da6daf5704805b9b5810f4575de7016b8f6cb1da1d7bb8a -size 200 diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap b/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap deleted file mode 100644 index c9416940..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8fca8ddbda3f7bb2f6e7553297c18f3ab8f8b73d64b5c9f81a3695ad9379d403 -size 1058400 diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json b/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json deleted file mode 100644 index b13b8ec9..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json +++ /dev/null @@ -1 +0,0 @@ -{"image": {"device": "cpu", "shape": [50, 3, 84, 84], "dtype": "torch.uint8"}, "state": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap b/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap deleted file mode 100644 index 3bae16df..00000000 --- a/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7b3e3e12896d553c208ee152f6d447c877c435e15d010c4a6171966d5b8a0c0b -size 800 diff --git a/tests/data/xarm_lift_medium/stats.pth b/tests/data/xarm_lift_medium/stats.pth index 0accffb0..3ab4e05b 100644 Binary files a/tests/data/xarm_lift_medium/stats.pth and b/tests/data/xarm_lift_medium/stats.pth differ diff --git a/tests/scripts/mock_dataset.py b/tests/scripts/mock_dataset.py index 044417aa..72480666 100644 --- a/tests/scripts/mock_dataset.py +++ b/tests/scripts/mock_dataset.py @@ -26,16 +26,23 @@ import torch def mock_dataset(in_data_dir, out_data_dir, num_frames): in_data_dir = Path(in_data_dir) out_data_dir = Path(out_data_dir) + out_data_dir.mkdir(exist_ok=True, parents=True) # copy the first `n` frames for each data key so that we have real data in_data_dict = torch.load(in_data_dir / "data_dict.pth") out_data_dict = {key: in_data_dict[key][:num_frames].clone() for key in in_data_dict} torch.save(out_data_dict, out_data_dir / "data_dict.pth") - # copy the full mapping between data_id and episode since it's small - in_ids_per_ep_path = in_data_dir / "data_ids_per_episode.pth" - out_ids_per_ep_path = out_data_dir / "data_ids_per_episode.pth" - shutil.copy(in_ids_per_ep_path, out_ids_per_ep_path) + # recreate data_ids_per_episode that corresponds to the subset + episodes = in_data_dict["episode"][:num_frames].tolist() + data_ids_per_episode = {} + for idx, ep_id in enumerate(episodes): + if ep_id not in data_ids_per_episode: + data_ids_per_episode[ep_id] = [] + data_ids_per_episode[ep_id].append(idx) + for ep_id in data_ids_per_episode: + data_ids_per_episode[ep_id] = torch.tensor(data_ids_per_episode[ep_id]) + torch.save(data_ids_per_episode, out_data_dir / "data_ids_per_episode.pth") # copy the full statistics of dataset since it's small in_stats_path = in_data_dir / "stats.pth" diff --git a/tests/test_datasets.py b/tests/test_datasets.py index e5ca0099..e7777c16 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -9,43 +9,76 @@ from .utils import DEVICE, DEFAULT_CONFIG_PATH @pytest.mark.parametrize( - "env_name,dataset_id", + "env_name,dataset_id,policy_name", [ - ("simxarm", "lift"), - ("pusht", "pusht"), - ("aloha", "sim_insertion_human"), - ("aloha", "sim_insertion_scripted"), - ("aloha", "sim_transfer_cube_human"), - ("aloha", "sim_transfer_cube_scripted"), + ("simxarm", "xarm_lift_medium", "tdmpc"), + ("pusht", "pusht", "diffusion"), + ("aloha", "aloha_sim_insertion_human", "act"), + ("aloha", "aloha_sim_insertion_scripted", "act"), + ("aloha", "aloha_sim_transfer_cube_human", "act"), + ("aloha", "aloha_sim_transfer_cube_scripted", "act"), ], ) -def test_factory(env_name, dataset_id): +def test_factory(env_name, dataset_id, policy_name): cfg = init_hydra_config( DEFAULT_CONFIG_PATH, - overrides=[f"env={env_name}", f"env.task={dataset_id}", f"device={DEVICE}"] + overrides=[f"env={env_name}", f"dataset_id={dataset_id}", f"policy={policy_name}", f"device={DEVICE}"] ) dataset = make_dataset(cfg) + delta_timestamps = dataset.delta_timestamps + image_keys = dataset.image_keys item = dataset[0] - assert "action" in item - assert "episode" in item - assert "frame_id" in item - assert "timestamp" in item - assert "next.done" in item - # TODO(rcadene): should we rename it agent_pos? - assert "observation.state" in item - for key in dataset.image_keys: - img = item.get(key) - assert img.dtype == torch.float32 - # TODO(rcadene): we assume for now that image normalization takes place in the model - assert img.max() <= 1.0 - assert img.min() >= 0.0 + keys_ndim_required = [ + ("action", 1, True), + ("episode", 0, True), + ("frame_id", 0, True), + ("timestamp", 0, True), + # TODO(rcadene): should we rename it agent_pos? + ("observation.state", 1, True), + ("next.reward", 0, False), + ("next.done", 0, False), + ] - if "next.reward" not in item: - logging.warning(f'Missing "next.reward" key in dataset {dataset}.') - if "next.done" not in item: - logging.warning(f'Missing "next.done" key in dataset {dataset}.') + for key in image_keys: + keys_ndim_required.append( + (key, 3, True), + ) + + # test number of dimensions + for key, ndim, required in keys_ndim_required: + if key not in item: + if required: + assert key in item, f"{key}" + else: + logging.warning(f'Missing key in dataset: "{key}" not in {dataset}.') + continue + + if delta_timestamps is not None and key in delta_timestamps: + assert item[key].ndim == ndim + 1, f"{key}" + assert item[key].shape[0] == len(delta_timestamps[key]), f"{key}" + else: + assert item[key].ndim == ndim, f"{key}" + + if key in image_keys: + assert item[key].dtype == torch.float32, f"{key}" + # TODO(rcadene): we assume for now that image normalization takes place in the model + assert item[key].max() <= 1.0, f"{key}" + assert item[key].min() >= 0.0, f"{key}" + + if delta_timestamps is not None and key in delta_timestamps: + # test t,c,h,w + assert item[key].shape[1] == 3, f"{key}" + else: + # test c,h,w + assert item[key].shape[0] == 3, f"{key}" + + + if delta_timestamps is not None: + # test missing keys in delta_timestamps + for key in delta_timestamps: + assert key in item, f"{key}" # def test_compute_stats(): diff --git a/tests/test_policies.py b/tests/test_policies.py index e1b3a4b6..c79bff94 100644 --- a/tests/test_policies.py +++ b/tests/test_policies.py @@ -13,7 +13,7 @@ from .utils import DEVICE, DEFAULT_CONFIG_PATH "env_name,policy_name,extra_overrides", [ ("simxarm", "tdmpc", ["policy.mpc=true"]), - #("pusht", "tdmpc", ["policy.mpc=false"]), + ("pusht", "tdmpc", ["policy.mpc=false"]), ("pusht", "diffusion", []), # ("aloha", "act", ["env.task=sim_insertion", "dataset_id=aloha_sim_insertion_human"]), #("aloha", "act", ["env.task=sim_insertion", "dataset_id=aloha_sim_insertion_scripted"]),