Add pusht test artifact

.gitattributes

@@ -0,0 +1 @@
+*.memmap filter=lfs diff=lfs merge=lfs -text

.github/workflows/test.yml

@@ -17,6 +17,7 @@ jobs:
     runs-on: ubuntu-latest
     env:
       POETRY_VERSION: 1.8.1
+      DATA_DIR: tests/data
     steps:
       #----------------------------------------------
       #       check-out repo and set-up python

.gitignore

@@ -54,6 +54,7 @@ pip-log.txt
 pip-delete-this-directory.txt
 
 # Unit test / coverage reports
+!tests/data
 htmlcov/
 .tox/
 .nox/

lerobot/common/datasets/pusht.py

@@ -125,6 +125,9 @@ class PushtExperienceReplay(AbstractExperienceReplay):
         episode_ids = torch.from_numpy(dataset_dict.get_episode_idxs())
         num_episodes = dataset_dict.meta["episode_ends"].shape[0]
         total_frames = dataset_dict["action"].shape[0]
+        # to create test artifact
+        # num_episodes = 1
+        # total_frames = 50
         assert len(
             {dataset_dict[key].shape[0] for key in dataset_dict.keys()}  # noqa: SIM118
         ), "Some data type dont have the same number of total frames."
@@ -142,6 +145,8 @@ class PushtExperienceReplay(AbstractExperienceReplay):
         idxtd = 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
 

tests/data/meta.json

@@ -0,0 +1 @@
+{"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": "<class 'tensordict._td.TensorDict'>"}

tests/data/next/meta.json

@@ -0,0 +1 @@
+{"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": "<class 'tensordict._td.TensorDict'>"}

tests/data/next/observation/meta.json

@@ -0,0 +1 @@
+{"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": "<class 'tensordict._td.TensorDict'>"}

tests/data/observation/meta.json

@@ -0,0 +1 @@
+{"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": "<class 'tensordict._td.TensorDict'>"}