Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Feature] Roboset datasets #1743

Merged
merged 4 commits into from
Dec 13, 2023
Merged

[Feature] Roboset datasets #1743

merged 4 commits into from
Dec 13, 2023

Conversation

vmoens
Copy link
Contributor

@vmoens vmoens commented Dec 12, 2023 •
edited
Loading

cc @vikashplus

Example usage:

from torchrl.data.datasets.roboset import RobosetExperienceReplay
d = RobosetExperienceReplay("FK1-v4(expert)/FK1_MicroOpenRandom_v2d-v4", batch_size=32)

for batch in d:
    break

# data is organised by seed and episode, but stored contiguously
print(batch["seed"])
print(batch["episode"])

print(batch)

Prints

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False),
        done: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
        env_infos: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
                obs_dict: TensorDict(
                    fields={
                        approach_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        ee_pose: Tensor(shape=torch.Size([32, 6]), device=cpu, dtype=torch.float64, is_shared=False),
                        end_effector: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        goal_err: Tensor(shape=torch.Size([32, 15]), device=cpu, dtype=torch.float64, is_shared=False),
                        kettle_site0_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        knob1_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        knob2_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        knob3_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        knob4_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        leftdoor_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        light_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        microhandle_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        obj_goal: Tensor(shape=torch.Size([32, 15]), device=cpu, dtype=torch.float64, is_shared=False),
                        objs_jnt: Tensor(shape=torch.Size([32, 15]), device=cpu, dtype=torch.float64, is_shared=False),
                        objs_vel: Tensor(shape=torch.Size([32, 15]), device=cpu, dtype=torch.float64, is_shared=False),
                        pose_err: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False),
                        qpos: Tensor(shape=torch.Size([32, 29]), device=cpu, dtype=torch.float64, is_shared=False),
                        rightdoor_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        robot_jnt: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False),
                        robot_vel: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False),
                        slide_site_err: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        time: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False)},
                    batch_size=torch.Size([32]),
                    device=cpu,
                    is_shared=False),
                proprio_dict: TensorDict(
                    fields={
                        ee_pose: Tensor(shape=torch.Size([32, 6]), device=cpu, dtype=torch.float64, is_shared=False),
                        robot_jnt: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False),
                        robot_vel: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False)},
                    batch_size=torch.Size([32]),
                    device=cpu,
                    is_shared=False),
                rwd_dense: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                rwd_dict: TensorDict(
                    fields={
                        approach: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                        bonus: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                        dense: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                        done: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
                        obj_goal: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                        pose: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                        solved: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
                        sparse: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False)},
                    batch_size=torch.Size([32]),
                    device=cpu,
                    is_shared=False),
                rwd_sparse: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                solved: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
                state: TensorDict(
                    fields={
                        act: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float16, is_shared=False),
                        body_pos: Tensor(shape=torch.Size([32, 44, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        body_quat: Tensor(shape=torch.Size([32, 44, 4]), device=cpu, dtype=torch.float64, is_shared=False),
                        mocap_pos: Tensor(shape=torch.Size([32, 1, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        mocap_quat: Tensor(shape=torch.Size([32, 1, 4]), device=cpu, dtype=torch.float64, is_shared=False),
                        qpos: Tensor(shape=torch.Size([32, 29]), device=cpu, dtype=torch.float64, is_shared=False),
                        qvel: Tensor(shape=torch.Size([32, 29]), device=cpu, dtype=torch.float64, is_shared=False),
                        site_pos: Tensor(shape=torch.Size([32, 14, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                        site_quat: Tensor(shape=torch.Size([32, 14, 4]), device=cpu, dtype=torch.float64, is_shared=False),
                        time: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False)},
                    batch_size=torch.Size([32]),
                    device=cpu,
                    is_shared=False),
                time: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False),
                visual_dict: TensorDict(
                    fields={
                        rgb:Franka_wrist_cam:256x256:2d: Tensor(shape=torch.Size([32, 256, 256, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
                        rgb:left_cam:256x256:2d: Tensor(shape=torch.Size([32, 256, 256, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
                        rgb:right_cam:256x256:2d: Tensor(shape=torch.Size([32, 256, 256, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
                        rgb:top_cam:256x256:2d: Tensor(shape=torch.Size([32, 256, 256, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
                        time: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float64, is_shared=False)},
                    batch_size=torch.Size([32]),
                    device=cpu,
                    is_shared=False)},
            batch_size=torch.Size([32]),
            device=cpu,
            is_shared=False),
        episode: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
        index: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
                env_infos: TensorDict(
                    ..., # same as above
                    batch_size=torch.Size([32]),
                    device=cpu,
                    is_shared=False),
                observation: Tensor(shape=torch.Size([32, 75]), device=cpu, dtype=torch.float64, is_shared=False),
                reward: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float64, is_shared=False),
                terminated: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([32]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([32, 75]), device=cpu, dtype=torch.float64, is_shared=False),
        seed: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
        time: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([32]),
    device=cpu,
    is_shared=False)

TODO:

  • Allow to discard the info dict if not needed
  • Tests

Extra features:

  • Build the "next" entry dynamically during sampling
  • Sample entire trajectories
  • Sample contiguous data (eg, sample 10 sequences of 25 trials)

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Dec 12, 2023 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/1743

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (18 Unrelated Failures)

As of commit 5f3b12c with merge base 3885b78 (image):

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Dec 12, 2023
@vmoens vmoens added the enhancement New feature or request label Dec 12, 2023
@github-actions GitHub Actions
Copy link

github-actions bot commented Dec 12, 2023 •
edited
Loading

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 92. Improved: $\large\color{#35bf28}3$. Worsened: $\large\color{#d91a1a}4$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_single 0.1242s 0.1229s 8.1363 Ops/s 8.0031 Ops/s $\color{#35bf28}+1.66\%$
test_sync 0.1025s 0.1020s 9.8005 Ops/s 9.7567 Ops/s $\color{#35bf28}+0.45\%$
test_async 0.2685s 99.7953ms 10.0205 Ops/s 9.9564 Ops/s $\color{#35bf28}+0.64\%$
test_single_pixels 0.1333s 0.1330s 7.5183 Ops/s 7.3696 Ops/s $\color{#35bf28}+2.02\%$
test_sync_pixels 96.7335ms 95.6047ms 10.4597 Ops/s 10.5956 Ops/s $\color{#d91a1a}-1.28\%$
test_async_pixels 0.2513s 92.4829ms 10.8128 Ops/s 10.9303 Ops/s $\color{#d91a1a}-1.07\%$
test_simple 0.9728s 0.9128s 1.0956 Ops/s 1.0870 Ops/s $\color{#35bf28}+0.79\%$
test_transformed 1.2342s 1.1593s 0.8626 Ops/s 0.8665 Ops/s $\color{#d91a1a}-0.45\%$
test_serial 2.6631s 2.5914s 0.3859 Ops/s 0.3888 Ops/s $\color{#d91a1a}-0.74\%$
test_parallel 2.5838s 2.5075s 0.3988 Ops/s 0.3941 Ops/s $\color{#35bf28}+1.19\%$
test_step_mdp_speed[True-True-True-True-True] 0.1034ms 35.3208μs 28.3119 KOps/s 30.1443 KOps/s $\textbf{\color{#d91a1a}-6.08\%}$
test_step_mdp_speed[True-True-True-True-False] 36.3720μs 20.1400μs 49.6525 KOps/s 51.3193 KOps/s $\color{#d91a1a}-3.25\%$
test_step_mdp_speed[True-True-True-False-True] 37.6630μs 19.6018μs 51.0157 KOps/s 53.4766 KOps/s $\color{#d91a1a}-4.60\%$
test_step_mdp_speed[True-True-True-False-False] 36.4520μs 11.3655μs 87.9854 KOps/s 89.1263 KOps/s $\color{#d91a1a}-1.28\%$
test_step_mdp_speed[True-True-False-True-True] 64.9840μs 36.6180μs 27.3089 KOps/s 28.8037 KOps/s $\textbf{\color{#d91a1a}-5.19\%}$
test_step_mdp_speed[True-True-False-True-False] 49.4030μs 21.8474μs 45.7720 KOps/s 46.5610 KOps/s $\color{#d91a1a}-1.69\%$
test_step_mdp_speed[True-True-False-False-True] 43.2520μs 21.1246μs 47.3382 KOps/s 48.5078 KOps/s $\color{#d91a1a}-2.41\%$
test_step_mdp_speed[True-True-False-False-False] 34.9820μs 13.2428μs 75.5129 KOps/s 76.1691 KOps/s $\color{#d91a1a}-0.86\%$
test_step_mdp_speed[True-False-True-True-True] 58.8730μs 37.6106μs 26.5882 KOps/s 27.0588 KOps/s $\color{#d91a1a}-1.74\%$
test_step_mdp_speed[True-False-True-True-False] 48.9030μs 23.7451μs 42.1139 KOps/s 42.8666 KOps/s $\color{#d91a1a}-1.76\%$
test_step_mdp_speed[True-False-True-False-True] 38.7420μs 21.0755μs 47.4484 KOps/s 48.7305 KOps/s $\color{#d91a1a}-2.63\%$
test_step_mdp_speed[True-False-True-False-False] 33.2910μs 13.3510μs 74.9007 KOps/s 76.7236 KOps/s $\color{#d91a1a}-2.38\%$
test_step_mdp_speed[True-False-False-True-True] 64.8740μs 38.9097μs 25.7006 KOps/s 25.9089 KOps/s $\color{#d91a1a}-0.80\%$
test_step_mdp_speed[True-False-False-True-False] 62.2730μs 25.7085μs 38.8976 KOps/s 39.5192 KOps/s $\color{#d91a1a}-1.57\%$
test_step_mdp_speed[True-False-False-False-True] 48.7930μs 22.8584μs 43.7476 KOps/s 44.4227 KOps/s $\color{#d91a1a}-1.52\%$
test_step_mdp_speed[True-False-False-False-False] 39.4520μs 15.2884μs 65.4092 KOps/s 67.0090 KOps/s $\color{#d91a1a}-2.39\%$
test_step_mdp_speed[False-True-True-True-True] 61.0730μs 37.3685μs 26.7605 KOps/s 27.0584 KOps/s $\color{#d91a1a}-1.10\%$
test_step_mdp_speed[False-True-True-True-False] 41.3020μs 23.8793μs 41.8772 KOps/s 42.8134 KOps/s $\color{#d91a1a}-2.19\%$
test_step_mdp_speed[False-True-True-False-True] 41.7920μs 25.3960μs 39.3763 KOps/s 40.7719 KOps/s $\color{#d91a1a}-3.42\%$
test_step_mdp_speed[False-True-True-False-False] 39.7320μs 15.5186μs 64.4389 KOps/s 67.2456 KOps/s $\color{#d91a1a}-4.17\%$
test_step_mdp_speed[False-True-False-True-True] 66.5330μs 39.4415μs 25.3540 KOps/s 25.9361 KOps/s $\color{#d91a1a}-2.24\%$
test_step_mdp_speed[False-True-False-True-False] 48.1730μs 25.7682μs 38.8076 KOps/s 39.7048 KOps/s $\color{#d91a1a}-2.26\%$
test_step_mdp_speed[False-True-False-False-True] 48.1630μs 26.8001μs 37.3132 KOps/s 38.1340 KOps/s $\color{#d91a1a}-2.15\%$
test_step_mdp_speed[False-True-False-False-False] 44.7920μs 17.0547μs 58.6347 KOps/s 60.4294 KOps/s $\color{#d91a1a}-2.97\%$
test_step_mdp_speed[False-False-True-True-True] 61.9930μs 41.3483μs 24.1848 KOps/s 24.8378 KOps/s $\color{#d91a1a}-2.63\%$
test_step_mdp_speed[False-False-True-True-False] 53.5230μs 27.6473μs 36.1700 KOps/s 36.9896 KOps/s $\color{#d91a1a}-2.22\%$
test_step_mdp_speed[False-False-True-False-True] 73.9440μs 27.1874μs 36.7818 KOps/s 38.3769 KOps/s $\color{#d91a1a}-4.16\%$
test_step_mdp_speed[False-False-True-False-False] 37.5420μs 17.1182μs 58.4172 KOps/s 60.0562 KOps/s $\color{#d91a1a}-2.73\%$
test_step_mdp_speed[False-False-False-True-True] 65.4030μs 42.2647μs 23.6604 KOps/s 23.7288 KOps/s $\color{#d91a1a}-0.29\%$
test_step_mdp_speed[False-False-False-True-False] 52.3930μs 29.4139μs 33.9975 KOps/s 34.4956 KOps/s $\color{#d91a1a}-1.44\%$
test_step_mdp_speed[False-False-False-False-True] 59.9530μs 28.1581μs 35.5138 KOps/s 35.9578 KOps/s $\color{#d91a1a}-1.23\%$
test_step_mdp_speed[False-False-False-False-False] 45.7230μs 18.7106μs 53.4457 KOps/s 53.9225 KOps/s $\color{#d91a1a}-0.88\%$
test_values[generalized_advantage_estimate-True-True] 27.3975ms 26.8455ms 37.2502 Ops/s 36.8972 Ops/s $\color{#35bf28}+0.96\%$
test_values[vec_generalized_advantage_estimate-True-True] 92.0897ms 3.4285ms 291.6728 Ops/s 286.3273 Ops/s $\color{#35bf28}+1.87\%$
test_values[td0_return_estimate-False-False] 0.1071ms 65.5644μs 15.2522 KOps/s 14.6745 KOps/s $\color{#35bf28}+3.94\%$
test_values[td1_return_estimate-False-False] 59.9023ms 59.4674ms 16.8159 Ops/s 17.0224 Ops/s $\color{#d91a1a}-1.21\%$
test_values[vec_td1_return_estimate-False-False] 2.1445ms 1.7904ms 558.5485 Ops/s 553.0503 Ops/s $\color{#35bf28}+0.99\%$
test_values[td_lambda_return_estimate-True-False] 97.9230ms 95.9976ms 10.4169 Ops/s 10.5427 Ops/s $\color{#d91a1a}-1.19\%$
test_values[vec_td_lambda_return_estimate-True-False] 2.0959ms 1.8126ms 551.6906 Ops/s 554.6363 Ops/s $\color{#d91a1a}-0.53\%$
test_gae_speed[generalized_advantage_estimate-False-1-512] 26.1237ms 25.5769ms 39.0978 Ops/s 39.0614 Ops/s $\color{#35bf28}+0.09\%$
test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 0.8768ms 0.7226ms 1.3838 KOps/s 1.3357 KOps/s $\color{#35bf28}+3.61\%$
test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 0.7475ms 0.6938ms 1.4413 KOps/s 1.4339 KOps/s $\color{#35bf28}+0.52\%$
test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 1.5592ms 1.4884ms 671.8495 Ops/s 669.7237 Ops/s $\color{#35bf28}+0.32\%$
test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 0.9735ms 0.7030ms 1.4224 KOps/s 1.3893 KOps/s $\color{#35bf28}+2.38\%$
test_dqn_speed 7.6615ms 1.4947ms 669.0134 Ops/s 674.8446 Ops/s $\color{#d91a1a}-0.86\%$
test_ddpg_speed 4.3530ms 3.3306ms 300.2438 Ops/s 296.2869 Ops/s $\color{#35bf28}+1.34\%$
test_sac_speed 9.9824ms 9.3349ms 107.1253 Ops/s 106.6159 Ops/s $\color{#35bf28}+0.48\%$
test_redq_speed 17.0175ms 16.5430ms 60.4486 Ops/s 59.7677 Ops/s $\color{#35bf28}+1.14\%$
test_redq_deprec_speed 13.9366ms 13.0217ms 76.7950 Ops/s 76.0985 Ops/s $\color{#35bf28}+0.92\%$
test_td3_speed 18.2445ms 9.5534ms 104.6752 Ops/s 104.3148 Ops/s $\color{#35bf28}+0.35\%$
test_cql_speed 33.4905ms 31.7188ms 31.5270 Ops/s 30.8379 Ops/s $\color{#35bf28}+2.23\%$
test_a2c_speed 7.9251ms 7.1547ms 139.7684 Ops/s 141.4673 Ops/s $\color{#d91a1a}-1.20\%$
test_ppo_speed 8.4874ms 7.4763ms 133.7555 Ops/s 134.9571 Ops/s $\color{#d91a1a}-0.89\%$
test_reinforce_speed 7.3322ms 6.1658ms 162.1837 Ops/s 161.8827 Ops/s $\color{#35bf28}+0.19\%$
test_iql_speed 28.2021ms 27.1861ms 36.7835 Ops/s 37.0318 Ops/s $\color{#d91a1a}-0.67\%$
test_sample_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 3.0867ms 2.4918ms 401.3136 Ops/s 398.6703 Ops/s $\color{#35bf28}+0.66\%$
test_sample_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 3.6094ms 2.6735ms 374.0380 Ops/s 330.5130 Ops/s $\textbf{\color{#35bf28}+13.17\%}$
test_sample_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 4.3275ms 2.6798ms 373.1595 Ops/s 374.6378 Ops/s $\color{#d91a1a}-0.39\%$
test_sample_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 3.1948ms 2.5066ms 398.9418 Ops/s 398.7815 Ops/s $\color{#35bf28}+0.04\%$
test_sample_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3.6272ms 2.6785ms 373.3384 Ops/s 372.1156 Ops/s $\color{#35bf28}+0.33\%$
test_sample_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3.9342ms 2.6856ms 372.3588 Ops/s 371.8172 Ops/s $\color{#35bf28}+0.15\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 3.1920ms 2.5061ms 399.0239 Ops/s 397.1567 Ops/s $\color{#35bf28}+0.47\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 3.6760ms 2.6905ms 371.6774 Ops/s 371.1622 Ops/s $\color{#35bf28}+0.14\%$
test_sample_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 4.3357ms 2.6875ms 372.0956 Ops/s 372.4530 Ops/s $\color{#d91a1a}-0.10\%$
test_iterate_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 2.7052ms 2.4764ms 403.8104 Ops/s 397.6308 Ops/s $\color{#35bf28}+1.55\%$
test_iterate_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 4.1167ms 2.6881ms 372.0097 Ops/s 370.8348 Ops/s $\color{#35bf28}+0.32\%$
test_iterate_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3.5354ms 2.6896ms 371.8070 Ops/s 371.0242 Ops/s $\color{#35bf28}+0.21\%$
test_iterate_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 2.8771ms 2.5070ms 398.8880 Ops/s 397.9348 Ops/s $\color{#35bf28}+0.24\%$
test_iterate_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3.9019ms 2.6830ms 372.7114 Ops/s 371.1457 Ops/s $\color{#35bf28}+0.42\%$
test_iterate_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 4.0641ms 2.7116ms 368.7904 Ops/s 368.1181 Ops/s $\color{#35bf28}+0.18\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 2.9521ms 2.5211ms 396.6496 Ops/s 400.6929 Ops/s $\color{#d91a1a}-1.01\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 4.1880ms 2.6976ms 370.6967 Ops/s 369.4103 Ops/s $\color{#35bf28}+0.35\%$
test_iterate_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 4.0274ms 2.6787ms 373.3131 Ops/s 370.5437 Ops/s $\color{#35bf28}+0.75\%$
test_populate_rb[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 0.1949s 18.9536ms 52.7604 Ops/s 53.8975 Ops/s $\color{#d91a1a}-2.11\%$
test_populate_rb[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 0.1190s 15.1513ms 66.0010 Ops/s 58.0190 Ops/s $\textbf{\color{#35bf28}+13.76\%}$
test_populate_rb[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 0.1191s 17.2981ms 57.8100 Ops/s 58.1102 Ops/s $\color{#d91a1a}-0.52\%$
test_populate_rb[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 0.1194s 17.2570ms 57.9476 Ops/s 66.7398 Ops/s $\textbf{\color{#d91a1a}-13.17\%}$
test_populate_rb[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 0.1200s 17.2269ms 58.0486 Ops/s 66.3505 Ops/s $\textbf{\color{#d91a1a}-12.51\%}$
test_populate_rb[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 0.1192s 17.2850ms 57.8537 Ops/s 58.1492 Ops/s $\color{#d91a1a}-0.51\%$
test_populate_rb[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 0.1192s 17.3308ms 57.7006 Ops/s 57.9377 Ops/s $\color{#d91a1a}-0.41\%$
test_populate_rb[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 0.1193s 15.1450ms 66.0283 Ops/s 58.2095 Ops/s $\textbf{\color{#35bf28}+13.43\%}$
test_populate_rb[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 0.1192s 17.2773ms 57.8793 Ops/s 58.1788 Ops/s $\color{#d91a1a}-0.51\%$

@vmoens
Copy link
Contributor Author

vmoens commented Dec 12, 2023

@vikashplus the class now downloads data from hf!
test:

from torchrl.data.datasets.roboset import RobosetExperienceReplay
from tempfile import TemporaryDirectory
for data in RobosetExperienceReplay.available_datasets:
    print("\n\n\ndata",  data)
    with TemporaryDirectory() as tmpdir:
        d = RobosetExperienceReplay(data, batch_size=32, root=tmpdir)
        for batch in d:
            print(batch)
            break

@vmoens vmoens added the Data Data-related PR, will launch data-related jobs label Dec 12, 2023
@vmoens vmoens marked this pull request as ready for review December 13, 2023 18:43
@vmoens vmoens merged commit 0906206 into main Dec 13, 2023
45 of 63 checks passed
@vmoens vmoens mentioned this pull request Dec 20, 2023
Open
36 tasks
@vmoens vmoens deleted the roboset branch February 27, 2024 00:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Data Data-related PR, will launch data-related jobs enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@vmoens @facebook-github-bot