amend

pytorch · vmoens · Mar 16, 2023 · Feb 20, 2023 · Feb 20, 2023 · Feb 21, 2023
commit 71e84d25ae906012667ca2ebc3c3dfe011544eb7
diff --git a/.circleci/unittest/linux_libs/scripts_d4rl/run_test.sh b/.circleci/unittest/linux_libs/scripts_d4rl/run_test.sh
@@ -35,11 +35,7 @@ conda env config vars set LD_PRELOAD=$LD_PRELOAD:$STDC_LOC
 conda deactivate && conda activate ./env
 
 # this workflow only tests the libs
-python -c "import habitat;import habitat.gym"
-python -c """from torchrl.envs.libs.habitat import HabitatEnv
-env = HabitatEnv('HabitatRenderPick-v0')
-env.reset()
-"""
+python -c "import gym, d4rl"
 
 python .circleci/unittest/helpers/coverage_run_parallel.py -m pytest test/test_libs.py --instafail -v --durations 20 --capture no -k TestD4RL
 coverage combine

diff --git a/test/test_libs.py b/test/test_libs.py
@@ -23,6 +23,7 @@
 from torchrl.collectors import MultiaSyncDataCollector
 from torchrl.collectors.collectors import RandomPolicy
 from torchrl.data.datasets.d4rl import _has_d4rl, D4RL_ERR, D4RLExperienceReplay
+from torchrl.data.replay_buffers import SamplerWithoutReplacement
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.brax import _has_brax, BraxEnv
 from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv, DMControlWrapper
@@ -1163,6 +1164,22 @@ def test_terminate_on_end(self, task):
             data_from_env._storage._storage.select(*keys),
         )
 
+    @pytest.mark.parametrize("split_trajs", [True, False])
+    def test_d4rl_iteration(self, task, split_trajs):
+        batch_size = 3
+        data = D4RLExperienceReplay(
+            task,
+            split_trajs=split_trajs,
+            from_env=False,
+            terminate_on_end=True,
+            batch_size=batch_size,
+            sampler=SamplerWithoutReplacement(drop_last=True),
+        )
+        i = 0
+        for sample in data:  # noqa: B007
+            i += 1
+        assert len(data) // i == batch_size
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()