torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/data/datasets/d4rl_infos.py

@@ -0,0 +1,187 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+D4RL_DATASETS = {
+    "maze2d-open-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-open-sparse.hdf5",
+    "maze2d-umaze-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-umaze-sparse-v1.hdf5",
+    "maze2d-medium-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-medium-sparse-v1.hdf5",
+    "maze2d-large-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-large-sparse-v1.hdf5",
+    "maze2d-eval-umaze-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-eval-umaze-sparse-v1.hdf5",
+    "maze2d-eval-medium-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-eval-medium-sparse-v1.hdf5",
+    "maze2d-eval-large-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-eval-large-sparse-v1.hdf5",
+    "maze2d-open-dense-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-open-dense.hdf5",
+    "maze2d-umaze-dense-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-umaze-dense-v1.hdf5",
+    "maze2d-medium-dense-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-medium-dense-v1.hdf5",
+    "maze2d-large-dense-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-large-dense-v1.hdf5",
+    "maze2d-eval-umaze-dense-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-eval-umaze-dense-v1.hdf5",
+    "maze2d-eval-medium-dense-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-eval-medium-dense-v1.hdf5",
+    "maze2d-eval-large-dense-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-eval-large-dense-v1.hdf5",
+    "minigrid-fourrooms-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/minigrid/minigrid4rooms.hdf5",
+    "minigrid-fourrooms-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/minigrid/minigrid4rooms_random.hdf5",
+    "pen-human-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/pen-v0_demos_clipped.hdf5",
+    "pen-cloned-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/pen-demos-v0-bc-combined.hdf5",
+    "pen-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/pen-v0_expert_clipped.hdf5",
+    "hammer-human-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/hammer-v0_demos_clipped.hdf5",
+    "hammer-cloned-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/hammer-demos-v0-bc-combined.hdf5",
+    "hammer-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/hammer-v0_expert_clipped.hdf5",
+    "relocate-human-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/relocate-v0_demos_clipped.hdf5",
+    "relocate-cloned-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/relocate-demos-v0-bc-combined.hdf5",
+    "relocate-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/relocate-v0_expert_clipped.hdf5",
+    "door-human-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/door-v0_demos_clipped.hdf5",
+    "door-cloned-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/door-demos-v0-bc-combined.hdf5",
+    "door-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/door-v0_expert_clipped.hdf5",
+    "halfcheetah-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_random.hdf5",
+    "halfcheetah-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_medium.hdf5",
+    "halfcheetah-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_expert.hdf5",
+    "halfcheetah-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_mixed.hdf5",
+    "halfcheetah-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/halfcheetah_medium_expert.hdf5",
+    "walker2d-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_random.hdf5",
+    "walker2d-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_medium.hdf5",
+    "walker2d-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_expert.hdf5",
+    "walker2d-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker_mixed.hdf5",
+    "walker2d-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/walker2d_medium_expert.hdf5",
+    "hopper-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_random.hdf5",
+    "hopper-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_medium.hdf5",
+    "hopper-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_expert.hdf5",
+    "hopper-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_mixed.hdf5",
+    "hopper-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/hopper_medium_expert.hdf5",
+    "ant-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_random.hdf5",
+    "ant-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_medium.hdf5",
+    "ant-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_expert.hdf5",
+    "ant-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_mixed.hdf5",
+    "ant-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_medium_expert.hdf5",
+    "ant-random-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco/ant_random_expert.hdf5",
+    "antmaze-umaze-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_u-maze_noisy_multistart_False_multigoal_False_sparse.hdf5",
+    "antmaze-umaze-diverse-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_u-maze_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-medium-play-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_big-maze_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-medium-diverse-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_big-maze_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-large-play-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-large-diverse-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-umaze-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_u-maze_noisy_multistart_False_multigoal_False_sparse_fixed.hdf5",
+    "antmaze-umaze-diverse-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_u-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5",
+    "antmaze-medium-play-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_big-maze_noisy_multistart_True_multigoal_False_sparse_fixed.hdf5",
+    "antmaze-medium-diverse-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_big-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5",
+    "antmaze-large-play-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_False_sparse_fixed.hdf5",
+    "antmaze-large-diverse-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5",
+    "flow-ring-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/flow/flow-ring-v0-random.hdf5",
+    "flow-ring-controller-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/flow/flow-ring-v0-idm.hdf5",
+    "flow-merge-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/flow/flow-merge-v0-random.hdf5",
+    "flow-merge-controller-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/flow/flow-merge-v0-idm.hdf5",
+    "kitchen-complete-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/kitchen/mini_kitchen_microwave_kettle_light_slider-v0.hdf5",
+    "kitchen-partial-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/kitchen/kitchen_microwave_kettle_light_slider-v0.hdf5",
+    "kitchen-mixed-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/kitchen/kitchen_microwave_kettle_bottomburner_light-v0.hdf5",
+    "carla-lane-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/carla/carla_lane_follow_flat-v0.hdf5",
+    "carla-town-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/carla/carla_town_subsamp_flat-v0.hdf5",
+    "carla-town-full-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/carla/carla_town_flat-v0.hdf5",
+    "bullet-halfcheetah-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-halfcheetah_random.hdf5",
+    "bullet-halfcheetah-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-halfcheetah_medium.hdf5",
+    "bullet-halfcheetah-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-halfcheetah_expert.hdf5",
+    "bullet-halfcheetah-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-halfcheetah_medium_expert.hdf5",
+    "bullet-halfcheetah-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-halfcheetah_medium_replay.hdf5",
+    "bullet-hopper-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-hopper_random.hdf5",
+    "bullet-hopper-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-hopper_medium.hdf5",
+    "bullet-hopper-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-hopper_expert.hdf5",
+    "bullet-hopper-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-hopper_medium_expert.hdf5",
+    "bullet-hopper-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-hopper_medium_replay.hdf5",
+    "bullet-ant-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-ant_random.hdf5",
+    "bullet-ant-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-ant_medium.hdf5",
+    "bullet-ant-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-ant_expert.hdf5",
+    "bullet-ant-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-ant_medium_expert.hdf5",
+    "bullet-ant-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-ant_medium_replay.hdf5",
+    "bullet-walker2d-random-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-walker2d_random.hdf5",
+    "bullet-walker2d-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-walker2d_medium.hdf5",
+    "bullet-walker2d-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-walker2d_expert.hdf5",
+    "bullet-walker2d-medium-expert-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-walker2d_medium_expert.hdf5",
+    "bullet-walker2d-medium-replay-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-walker2d_medium_replay.hdf5",
+    "bullet-maze2d-open-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-maze2d-open-sparse.hdf5",
+    "bullet-maze2d-umaze-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-maze2d-umaze-sparse.hdf5",
+    "bullet-maze2d-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-maze2d-medium-sparse.hdf5",
+    "bullet-maze2d-large-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/bullet/bullet-maze2d-large-sparse.hdf5",
+    "halfcheetah-random-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/halfcheetah_random-v1.hdf5",
+    "halfcheetah-random-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/halfcheetah_random-v2.hdf5",
+    "halfcheetah-medium-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/halfcheetah_medium-v1.hdf5",
+    "halfcheetah-medium-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/halfcheetah_medium-v2.hdf5",
+    "halfcheetah-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/halfcheetah_expert-v1.hdf5",
+    "halfcheetah-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/halfcheetah_expert-v2.hdf5",
+    "halfcheetah-medium-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/halfcheetah_medium_replay-v1.hdf5",
+    "halfcheetah-medium-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/halfcheetah_medium_replay-v2.hdf5",
+    "halfcheetah-full-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/halfcheetah_full_replay-v1.hdf5",
+    "halfcheetah-full-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/halfcheetah_full_replay-v2.hdf5",
+    "halfcheetah-medium-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/halfcheetah_medium_expert-v1.hdf5",
+    "halfcheetah-medium-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/halfcheetah_medium_expert-v2.hdf5",
+    "hopper-random-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/hopper_random-v1.hdf5",
+    "hopper-random-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/hopper_random-v2.hdf5",
+    "hopper-medium-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/hopper_medium-v1.hdf5",
+    "hopper-medium-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/hopper_medium-v2.hdf5",
+    "hopper-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/hopper_expert-v1.hdf5",
+    "hopper-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/hopper_expert-v2.hdf5",
+    "hopper-medium-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/hopper_medium_replay-v1.hdf5",
+    "hopper-medium-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/hopper_medium_replay-v2.hdf5",
+    "hopper-full-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/hopper_full_replay-v1.hdf5",
+    "hopper-full-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/hopper_full_replay-v2.hdf5",
+    "hopper-medium-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/hopper_medium_expert-v1.hdf5",
+    "hopper-medium-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/hopper_medium_expert-v2.hdf5",
+    "walker2d-random-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/walker2d_random-v1.hdf5",
+    "walker2d-random-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/walker2d_random-v2.hdf5",
+    "walker2d-medium-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/walker2d_medium-v1.hdf5",
+    "walker2d-medium-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/walker2d_medium-v2.hdf5",
+    "walker2d-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/walker2d_expert-v1.hdf5",
+    "walker2d-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/walker2d_expert-v2.hdf5",
+    "walker2d-medium-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/walker2d_medium_replay-v1.hdf5",
+    "walker2d-medium-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/walker2d_medium_replay-v2.hdf5",
+    "walker2d-full-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/walker2d_full_replay-v1.hdf5",
+    "walker2d-full-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/walker2d_full_replay-v2.hdf5",
+    "walker2d-medium-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/walker2d_medium_expert-v1.hdf5",
+    "walker2d-medium-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/walker2d_medium_expert-v2.hdf5",
+    "ant-random-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/ant_random-v1.hdf5",
+    "ant-random-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/ant_random-v2.hdf5",
+    "ant-medium-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/ant_medium-v1.hdf5",
+    "ant-medium-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/ant_medium-v2.hdf5",
+    "ant-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/ant_expert-v1.hdf5",
+    "ant-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/ant_expert-v2.hdf5",
+    "ant-medium-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/ant_medium_replay-v1.hdf5",
+    "ant-medium-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/ant_medium_replay-v2.hdf5",
+    "ant-full-replay-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/ant_full_replay-v1.hdf5",
+    "ant-full-replay-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/ant_full_replay-v2.hdf5",
+    "ant-medium-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v1/ant_medium_expert-v1.hdf5",
+    "ant-medium-expert-v2": "http://rail.eecs.berkeley.edu/datasets/offline_rl/gym_mujoco_v2/ant_medium_expert-v2.hdf5",
+    "hammer-human-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/hammer-human-v1.hdf5",
+    "hammer-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/hammer-expert-v1.hdf5",
+    "hammer-cloned-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/hammer-cloned-v1.hdf5",
+    "pen-human-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/pen-human-v1.hdf5",
+    "pen-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/pen-expert-v1.hdf5",
+    "pen-cloned-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/pen-cloned-v1.hdf5",
+    "relocate-human-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/relocate-human-v1.hdf5",
+    "relocate-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/relocate-expert-v1.hdf5",
+    "relocate-cloned-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/relocate-cloned-v1.hdf5",
+    "door-human-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/door-human-v1.hdf5",
+    "door-expert-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/door-expert-v1.hdf5",
+    "door-cloned-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg_v1/door-cloned-v1.hdf5",
+    "antmaze-umaze-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v1/Ant_maze_umaze_noisy_multistart_False_multigoal_False_sparse.hdf5",
+    "antmaze-umaze-diverse-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v1/Ant_maze_umaze_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-medium-play-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v1/Ant_maze_medium_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-medium-diverse-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v1/Ant_maze_medium_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-large-diverse-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v1/Ant_maze_large_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-large-play-v1": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v1/Ant_maze_large_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-eval-umaze-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_umaze_eval_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-eval-umaze-diverse-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_umaze_eval_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-eval-medium-play-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_medium_eval_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "antmaze-eval-medium-diverse-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_medium_eval_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-eval-large-diverse-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_large_eval_noisy_multistart_True_multigoal_False_sparse.hdf5",
+    "antmaze-eval-large-play-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_new/Ant_maze_large_eval_noisy_multistart_True_multigoal_True_sparse.hdf5",
+    "door-human-longhorizon-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/door-v0_demos_clipped.hdf5",
+    "hammer-human-longhorizon-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/hammer-v0_demos_clipped.hdf5",
+    "pen-human-longhorizon-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/pen-v0_demos_clipped.hdf5",
+    "relocate-human-longhorizon-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/hand_dapg/relocate-v0_demos_clipped.hdf5",
+    "maze2d-umaze-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-umaze-sparse.hdf5",
+    "maze2d-medium-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-medium-sparse.hdf5",
+    "maze2d-large-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-large-sparse.hdf5",
+    "maze2d-umaze-dense-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-umaze-dense.hdf5",
+    "maze2d-medium-dense-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-medium-dense.hdf5",
+    "maze2d-large-dense-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/maze2d/maze2d-large-dense.hdf5",
+    "carla-lane-render-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/carla/carla_lane_follow-v0.hdf5",
+    "carla-town-render-v0": "http://rail.eecs.berkeley.edu/datasets/offline_rl/carla/carla_town_flat-v0.hdf5",
+}

torchrl/data/datasets/gen_dgrl.py

@@ -0,0 +1,375 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import importlib.util
+import os
+import shutil
+import tarfile
+import tempfile
+import typing as tp
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from tensordict import TensorDict
+from torchrl._utils import logger as torchrl_logger
+from torchrl.data.datasets.common import BaseDatasetExperienceReplay
+from torchrl.data.datasets.utils import _get_root_dir
+from torchrl.data.replay_buffers.storages import TensorStorage
+from torchrl.envs.utils import _classproperty
+
+_has_tqdm = importlib.util.find_spec("tqdm", None) is not None
+_has_requests = importlib.util.find_spec("requests", None) is not None
+
+
+class GenDGRLExperienceReplay(BaseDatasetExperienceReplay):
+    """Gen-DGRL Experience Replay dataset.
+
+    This dataset accompanies the paper "The Generalization Gap in Offline Reinforcement Learning".
+
+    Arxiv: https://arxiv.org/abs/2312.05742
+
+    GitHub: https://github.com/facebookresearch/gen_dgrl
+
+    The data format follows the :ref:`TED convention <TED-format>`.
+
+    This class gives you access to the ProcGen dataset. Each `dataset_id` registered
+    in `GenDGRLExperienceReplay.available_datasets` consists in a particular task
+    (`"bigfish"`, `"bossfight"`, ...) separated from a category (`"1M_E"`, `"1M_S"`, ...)
+    by a comma (`"bigfish-1M_E"`, ...).
+
+    During download and preparation, the data is downloaded as .tar files,
+    where each trajectory is stored independently in a .npy file. Each of these
+    files is extracted, written in a contiguous mmap tensor, and then cleared.
+    This process can take several minutes per dataset. On a cluster, it is advisable
+    to first run the download and preprocessing separately on different workers
+    or processes for different datasets, and launch the training script in a second time.
+
+    Args:
+        dataset_id (str): the dataset to be downloaded. Must be part of
+            :attr:`GenDGRLExperienceReplay.available_datasets`.
+        batch_size (int, optional): Batch-size used during sampling. Can be overridden by
+            `data.sample(batch_size)` if necessary.
+
+    Keyword Args:
+        root (Path or str, optional): The :class:`~torchrl.data.datasets.GenDGRLExperienceReplay`
+            dataset root directory.
+            The actual dataset memory-mapped files will be saved under
+            `<root>/<dataset_id>`. If none is provided, it defaults to
+            `~/.cache/torchrl/atari`.gen_dgrl`.
+        download (bool or str, optional): Whether the dataset should be downloaded if
+            not found. Defaults to ``True``. Download can also be passed as ``"force"``,
+            in which case the downloaded data will be overwritten.
+        sampler (Sampler, optional): the sampler to be used. If none is provided
+            a default RandomSampler() will be used.
+        writer (Writer, optional): the writer to be used. If none is provided
+            a default RoundRobinWriter() will be used.
+        collate_fn (callable, optional): merges a list of samples to form a
+            mini-batch of Tensor(s)/outputs.  Used when using batched
+            loading from a map-style dataset.
+        pin_memory (bool): whether pin_memory() should be called on the rb
+            samples.
+        prefetch (int, optional): number of next batches to be prefetched
+            using multithreading.
+        transform (Transform, optional): Transform to be executed when sample() is called.
+            To chain transforms use the :obj:`Compose` class.
+
+    Attributes:
+        available_datasets: a list of accepted entries to be downloaded. These
+            names correspond to the directory path in the huggingface dataset
+            repository. If possible, the list will be dynamically retrieved from
+            huggingface. If no internet connection is available, it a cached
+            version will be used.
+
+    Examples:
+        >>> import torch
+        >>> torch.manual_seed(0)
+        >>> from torchrl.data.datasets import GenDGRLExperienceReplay
+        >>> d = GenDGRLExperienceReplay("bigfish-1M_E", batch_size=32)
+        >>> for batch in d:
+        ...     break
+        >>> print(batch)
+
+    """
+
+    BASE_URL = "https://dl.fbaipublicfiles.com/DGRL/Procgen/Datasets/Compressed"
+    # number of files extracted at a time
+    _PROCESS_NPY_BATCH = 32
+    split_trajs: bool = False
+
+    @_classproperty
+    def available_datasets(cls):
+        datasets = [
+            "bigfish",
+            "bossfight",
+            "caveflyer",
+            "chaser",
+            "climber",
+            "coinrun",
+            "dodgeball",
+            "fruitbot",
+            "heist",
+            "jumper",
+            "leaper",
+            "maze",
+            "miner",
+            "ninja",
+            "plunder",
+            "starpilot",
+        ]
+        categories = [
+            "1M_E",
+            "1M_S",
+            "10M",
+            "25M",
+            "level_1_E",
+            "level_1_S",
+            "level_40_E",
+            "level_40_S",
+        ]
+
+        return ["-".join((ds, cat)) for cat in categories for ds in datasets]
+
+    def __init__(
+        self,
+        dataset_id: str,
+        batch_size: int | None = None,
+        *,
+        download: bool = True,
+        root: str | None = None,
+        **kwargs,
+    ):
+        self.dataset_id = dataset_id
+        try:
+            dataset, category = dataset_id.split("-")
+        except Exception:
+            category = dataset_id
+            dataset = None
+        self._dataset_name = dataset
+        self._category_name = category
+        if root is None:
+            root = _get_root_dir("gen_dgrl")
+            os.makedirs(root, exist_ok=True)
+        self.root = root
+        if download == "force" or (download and not self._is_downloaded()):
+            if download == "force" and os.path.exists(self.data_path_root):
+                shutil.rmtree(self.data_path_root)
+            storage = TensorStorage(self._download_and_preproc())
+        else:
+            storage = TensorStorage(TensorDict.load_memmap(self.data_path_root))
+        super().__init__(storage=storage, batch_size=batch_size, **kwargs)
+
+    @property
+    def data_path(self):
+        if self.split_trajs:
+            return Path(self.root) / (self.dataset_id + "_split")
+        return self.data_path_root
+
+    @property
+    def data_path_root(self):
+        return Path(self.root) / self.dataset_id
+
+    def _is_downloaded(self):
+        return os.path.exists(self.data_path_root)
+
+    def _download_and_preproc(self):
+        dataset, category = self._dataset_name, self._category_name
+        link = self._build_urls_with_category_name(dataset, category)
+        data_link = (category, self._fetch_file_name_from_link(link), link)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self._download_category_file(
+                tmpdir, skip_downloaded_files=True, link=data_link
+            )
+            return self._unpack_category_file(
+                tmpdir, clear_archive=True, link=data_link, category_name=category
+            )
+
+    @classmethod
+    def _build_urls_with_category_name(
+        cls, dataset, category_name: str
+    ) -> tp.List[str]:
+        path = [cls.BASE_URL, cls._convert_category_name(category_name)]
+        path += [f"{dataset}.tar.xz"]
+        return os.path.join(*path)
+
+    @staticmethod
+    def _convert_category_name(category_name: str) -> str:
+        if category_name == "1M_E":
+            return "1M/level_200/expert"
+        elif category_name == "1M_S":
+            return "1M/level_200/suboptimal"
+        elif category_name == "10M":
+            return "10M/level_200/expert"
+        elif category_name == "25M":
+            return "25M/level_200/expert"
+        elif category_name == "level_1_S":
+            return "100k/level_1/suboptimal"
+        elif category_name == "level_40_S":
+            return "100k/level_40/suboptimal"
+        elif category_name == "level_1_E":
+            return "100k/level_1/expert"
+        elif category_name == "level_40_E":
+            return "100k/level_40/expert"
+        else:
+            raise ValueError(f"Unrecognized category name {category_name}!")
+
+    @staticmethod
+    def _fetch_file_name_from_link(url: str) -> str:
+        return os.path.split(url)[-1]
+
+    @classmethod
+    def _get_category_len(cls, category_name):
+        if "1M" in category_name:
+            return 1_000_000
+        if "10M" in category_name:
+            return 10_000_000
+        if "25M" in category_name:
+            return 25_000_000
+        return 100_000
+
+    def _unpack_category_file(
+        self,
+        download_folder: str,
+        clear_archive: bool,
+        category_name,
+        link: str,
+        batch=None,
+    ):
+        if batch is None:
+            batch = self._PROCESS_NPY_BATCH
+        _, file_name, _ = link
+        file_path = os.path.join(download_folder, file_name)
+        torchrl_logger.info(
+            f"Unpacking dataset file {file_path} ({file_name}) to {download_folder}."
+        )
+        idx = 0
+        td_memmap = None
+        dataset_len = self._get_category_len(category_name)
+        if _has_tqdm:
+            from tqdm import tqdm
+
+            pbar = tqdm(total=dataset_len)
+        else:
+            pbar = None
+        mode = "r:xz" if str(file_path).endswith("xz") else "r"
+        full = False
+        with tarfile.open(file_path, mode) as tar:
+            members = list(tar.getmembers())
+            for i in range(0, len(members), batch):
+                if full:
+                    break
+                submembers = [
+                    member for member in members[i : i + batch] if member.isfile()
+                ]
+                for member in submembers:
+                    if pbar is not None:
+                        pbar.set_description(member.name)
+                    npybuffer = tar.extractfile(member=member)
+                    # npyfile = Path(download_folder) / member.name
+                    npfile = np.load(npybuffer, allow_pickle=True)
+                    td = TensorDict.from_dict(npfile.tolist(), auto_batch_size=True)
+                    td.set("observations", td.get("observations").to(torch.uint8))
+                    td.set(("next", "observation"), td.get("observations")[1:])
+                    td.set("observations", td.get("observations")[:-1])
+                    td.rename_key_("observations", "observation")
+                    td.rename_key_("dones", ("next", "done"))
+                    td.rename_key_("actions", "action")
+                    td.rename_key_("rewards", ("next", "reward"))
+                    td.set(
+                        ("next", "done"), td.get(("next", "done")).bool().unsqueeze(-1)
+                    )
+                    td.set(
+                        ("next", "truncated"),
+                        torch.zeros_like(td.get(("next", "done"))),
+                    )
+                    td.set(("next", "terminated"), td.get(("next", "done")))
+
+                    td.set(
+                        "terminated", torch.zeros_like(td.get(("next", "terminated")))
+                    )
+                    td.set("done", torch.zeros_like(td.get(("next", "done"))))
+                    td.set("truncated", torch.zeros_like(td.get(("next", "truncated"))))
+
+                    td.batch_size = td.get("observation").shape[:1]
+                    if td_memmap is None:
+                        td_memmap = (
+                            td[0]
+                            .expand(dataset_len)
+                            .memmap_like(self.data_path_root, num_threads=16)
+                        )
+                    idx_end = idx + td.shape[0]
+                    idx_end = min(idx_end, td_memmap.shape[0])
+                    if pbar is not None:
+                        pbar.update(td.shape[0])
+                    length = idx_end - idx
+                    if length > 0:
+                        if length != td.shape[0]:
+                            td_memmap[idx:idx_end] = td[:length]
+                        else:
+                            td_memmap[idx:idx_end] = td
+                    else:
+                        full = True
+                    idx = idx_end
+        if clear_archive:
+            os.remove(file_path)
+        return td_memmap
+
+    @classmethod
+    def _download_category_file(
+        cls,
+        download_folder: str,
+        skip_downloaded_files: bool,
+        link: str,
+    ):
+        _, file_name, url = link
+        file_path = os.path.join(download_folder, file_name)
+
+        if skip_downloaded_files and os.path.isfile(file_path):
+            torchrl_logger.info(f"Skipping {file_path}, already downloaded!")
+            return file_name, True
+
+        in_progress_folder = os.path.join(download_folder, "_in_progress")
+        os.makedirs(in_progress_folder, exist_ok=True)
+        in_progress_file_path = os.path.join(in_progress_folder, file_name)
+
+        torchrl_logger.info(
+            f"Downloading dataset file {file_name} ({url}) to {in_progress_file_path}."
+        )
+        cls._download_with_progress_bar(url, in_progress_file_path)
+
+        os.rename(in_progress_file_path, file_path)
+        return file_name, True
+
+    @classmethod
+    def _download_with_progress_bar(cls, url: str, file_path: str):
+        # taken from https://stackoverflow.com/a/62113293/986477
+        if not _has_requests:
+            raise ImportError(
+                "The requests package is required for Gen-DGRL dataset download."
+            )
+        import requests
+
+        resp = requests.get(url, stream=True)
+        total = int(resp.headers.get("content-length", 0))
+        if _has_tqdm:
+            from tqdm import tqdm
+
+            pbar = tqdm(
+                desc=file_path,
+                total=total,
+                unit="iB",
+                unit_scale=True,
+                unit_divisor=1024,
+            )
+        else:
+            pbar = None
+        with open(file_path, "wb") as file:
+            for data in resp.iter_content(chunk_size=1024):
+                size = file.write(data)
+                if pbar is not None:
+                    pbar.update(size)