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

[Feature] Roboset datasets #1743

Merged
merged 4 commits into from
Dec 13, 2023
Merged

[Feature] Roboset datasets #1743

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
410df84
init
vmoens Dec 12, 2023
caf1996
tests
vmoens Dec 12, 2023
83a6667
amend
vmoens Dec 12, 2023
5f3b12c
amend
vmoens Dec 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
amend
  • Loading branch information
@vmoens
vmoens committed Dec 12, 2023
commit 5f3b12cf7c6f7386f613ec08798012c6d39fe3c4
1 change: 1 addition & 0 deletions .github/unittest/linux_libs/scripts_roboset/environment.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,4 @@ dependencies:
- hydra-core
- huggingface_hub
- tqdm
- h5py
1 change: 1 addition & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,7 @@ def _main(argv):
"tqdm",
"scikit-learn",
"pandas",
"h5py",
],
"marl": ["vmas>=1.2.10", "pettingzoo>=1.24.1"],
}
Expand Down
6 changes: 6 additions & 0 deletions torchrl/data/datasets/roboset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@
from torchrl.data.replay_buffers.writers import Writer

_has_tqdm = importlib.util.find_spec("tqdm", None) is not None
_has_h5py = importlib.util.find_spec("h5py", None) is not None
_has_hf_hub = importlib.util.find_spec("huggingface_hub", None) is not None

_NAME_MATCH = KeyDependentDefaultDict(lambda key: key)
_NAME_MATCH["observations"] = "observation"
Expand Down Expand Up @@ -158,6 +160,10 @@ def __init__(
split_trajs: bool = False,
**env_kwargs,
):
if not _has_h5py or not _has_hf_hub:
raise ImportError(
"h5py and huggingface_hub are required for Roboset datasets."
)
if dataset_id not in self.available_datasets:
raise ValueError(
f"The dataset_id {dataset_id} isn't part of the accepted datasets. "
Expand Down
Loading