From bd1c80a7eb7db06d5ab60a71d244f360a3db2e91 Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 11 Jun 2026 14:52:32 +0200 Subject: [PATCH 01/11] isaac downloader --- .../src/molmo_spaces_isaac/downloader/main.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py b/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py index 7c9ba439..d72e2a98 100644 --- a/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py +++ b/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py @@ -5,7 +5,7 @@ from typing import Literal import tyro -from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, ResourceManager +from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, GCRemoteStorage, ResourceManager logger = logging.getLogger("molmospaces_resources") logger.setLevel(logging.DEBUG) @@ -84,8 +84,8 @@ class DownloadArgs: # If not provided, uses HF_TOKEN from environment hf_token: str | None = None - # Use R2 remote storage (HuggingFace by default) - use_r2: bool = False + # Storage to use (HuggingFace by default) + storage: list[Literal["hf", "r2", "gc"]] = field(default_factory=lambda: "hf") def main() -> int: @@ -112,14 +112,19 @@ def main() -> int: dataset_id ] + if args.storage == "hf": + remote_storage = HFRemoteStorage( + repo_id="allenai/molmospaces", + repo_prefix=TYPE_TO_PREFIX[args.type], + token=args.hf_token or os.getenv("HF_TOKEN"), + ) + elif args.storage == "r2": + remote_storage = R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") + else: + remote_storage = GCRemoteStorage(f"molmospaces-{TYPE_TO_PREFIX[args.type]}") + manager = ResourceManager( - remote_storage=R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") - if args.use_r2 - else HFRemoteStorage( - repo_id="allenai/molmospaces", - repo_prefix=TYPE_TO_PREFIX[args.type], - token=args.hf_token or os.getenv("HF_TOKEN"), - ), + remote_storage=remote_storage, data_type_to_source_to_version=sources_to_version, symlink_dir=args.install_dir, cache_dir=args.cache_dir / args.type, From 0050f5c52965e7c57cf32417fe6c5ad0bbda7320 Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 11 Jun 2026 14:57:12 +0200 Subject: [PATCH 02/11] maniskill downloader with gcs --- .../molmo_spaces_maniskill/downloader/main.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py b/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py index 7c9ba439..d72e2a98 100644 --- a/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py +++ b/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py @@ -5,7 +5,7 @@ from typing import Literal import tyro -from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, ResourceManager +from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, GCRemoteStorage, ResourceManager logger = logging.getLogger("molmospaces_resources") logger.setLevel(logging.DEBUG) @@ -84,8 +84,8 @@ class DownloadArgs: # If not provided, uses HF_TOKEN from environment hf_token: str | None = None - # Use R2 remote storage (HuggingFace by default) - use_r2: bool = False + # Storage to use (HuggingFace by default) + storage: list[Literal["hf", "r2", "gc"]] = field(default_factory=lambda: "hf") def main() -> int: @@ -112,14 +112,19 @@ def main() -> int: dataset_id ] + if args.storage == "hf": + remote_storage = HFRemoteStorage( + repo_id="allenai/molmospaces", + repo_prefix=TYPE_TO_PREFIX[args.type], + token=args.hf_token or os.getenv("HF_TOKEN"), + ) + elif args.storage == "r2": + remote_storage = R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") + else: + remote_storage = GCRemoteStorage(f"molmospaces-{TYPE_TO_PREFIX[args.type]}") + manager = ResourceManager( - remote_storage=R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") - if args.use_r2 - else HFRemoteStorage( - repo_id="allenai/molmospaces", - repo_prefix=TYPE_TO_PREFIX[args.type], - token=args.hf_token or os.getenv("HF_TOKEN"), - ), + remote_storage=remote_storage, data_type_to_source_to_version=sources_to_version, symlink_dir=args.install_dir, cache_dir=args.cache_dir / args.type, From 99050f2fa0bbfb8851742cf0907c28c710e881c5 Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 11 Jun 2026 15:03:34 +0200 Subject: [PATCH 03/11] molmospaces with gcs --- molmo_spaces/molmo_spaces_constants.py | 17 +++++++++++------ pyproject.toml | 2 +- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index 867dd56e..18fb83a9 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -13,11 +13,13 @@ from collections import defaultdict from pathlib import Path from copy import deepcopy +from typing import Literal import compress_json from molmospaces_resources import ( HFRemoteStorage, R2RemoteStorage, + GCRemoteStorage, ResourceManager, setup_resource_manager, str2bool, @@ -82,7 +84,7 @@ def resource_manager_log_level(log_level=logging.DEBUG): else None ) -USE_HUGGING_FACE = False # If True, HF_TOKEN needs to exist in the environment +STORAGE_IN_USE: Literal["gc", "hf", "r2"] = "gc" # If using hf, HF_TOKEN needs to exist in the environment DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ @@ -216,11 +218,14 @@ def register_user_grasp_library(root_name: str, path: Path, object_library: str) def _select_storage(): - return ( - HFRemoteStorage("allenai/molmospaces", repo_prefix="mujoco", token=os.getenv("HF_TOKEN")) - if USE_HUGGING_FACE - else R2RemoteStorage("mujoco-thor-resources") - ) + if STORAGE_IN_USE == "hf": + return HFRemoteStorage("allenai/molmospaces", repo_prefix="mujoco", token=os.getenv("HF_TOKEN")) + elif STORAGE_IN_USE == "r2": + return R2RemoteStorage("mujoco-thor-resources") + elif STORAGE_IN_USE == "gc": + return GCRemoteStorage("molmospaces-mujoco") + + raise NotImplementedError(f"Storage type {STORAGE_IN_USE} not supported") def get_resource_manager( diff --git a/pyproject.toml b/pyproject.toml index 68477b23..12e9cc71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ dependencies = [ "wandb>=0.18.10", "websockets", "pytest", # note (yejin): moved from dev deps since needed in Docker builds - "molmospaces-resources==0.0.1b4", + "molmospaces-resources==0.0.1", ] [project.optional-dependencies] dev = [ From 3127be5f388b1f5c0a57ffc157cad213d11bb5d3 Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 11 Jun 2026 15:06:36 +0200 Subject: [PATCH 04/11] usda_downloader with gcs --- scripts/assets/usda_downloader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/assets/usda_downloader.py b/scripts/assets/usda_downloader.py index 8bdb7841..d83879c1 100644 --- a/scripts/assets/usda_downloader.py +++ b/scripts/assets/usda_downloader.py @@ -4,7 +4,7 @@ from pxr import Sdf, Usd, Ar -from molmospaces_resources import setup_resource_manager, R2RemoteStorage +from molmospaces_resources import setup_resource_manager, GCRemoteStorage from molmo_spaces.molmo_spaces_constants import ASSETS_DIR, DATA_CACHE_DIR from molmospaces_resources import ResourceManager, SourceInfo @@ -36,7 +36,7 @@ def get_resource_manager() -> ResourceManager: global _RESOURCE_MANAGER if _RESOURCE_MANAGER is None: _RESOURCE_MANAGER = setup_resource_manager( - R2RemoteStorage("isaac-thor-resources"), + GCRemoteStorage("molmospaces-isaac"), symlink_dir=ISAAC_ASSETS_DIR, versions=ISAAC_DATA_TYPE_TO_SOURCE_TO_VERSION, cache_dir=ISAAC_DATA_CACHE_DIR, From 3ce1b89ca3ba7ad736fc326a50e21e6df057409d Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 11 Jun 2026 15:08:34 +0200 Subject: [PATCH 05/11] ruff --- molmo_spaces/molmo_spaces_constants.py | 8 ++++++-- .../src/molmo_spaces_isaac/downloader/main.py | 8 ++++---- .../src/molmo_spaces_maniskill/downloader/main.py | 8 ++++---- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index 18fb83a9..ec3b09b2 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -84,7 +84,9 @@ def resource_manager_log_level(log_level=logging.DEBUG): else None ) -STORAGE_IN_USE: Literal["gc", "hf", "r2"] = "gc" # If using hf, HF_TOKEN needs to exist in the environment +STORAGE_IN_USE: Literal["gc", "hf", "r2"] = ( + "gc" # If using hf, HF_TOKEN needs to exist in the environment +) DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ @@ -219,7 +221,9 @@ def register_user_grasp_library(root_name: str, path: Path, object_library: str) def _select_storage(): if STORAGE_IN_USE == "hf": - return HFRemoteStorage("allenai/molmospaces", repo_prefix="mujoco", token=os.getenv("HF_TOKEN")) + return HFRemoteStorage( + "allenai/molmospaces", repo_prefix="mujoco", token=os.getenv("HF_TOKEN") + ) elif STORAGE_IN_USE == "r2": return R2RemoteStorage("mujoco-thor-resources") elif STORAGE_IN_USE == "gc": diff --git a/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py b/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py index d72e2a98..ce989d9e 100644 --- a/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py +++ b/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py @@ -114,10 +114,10 @@ def main() -> int: if args.storage == "hf": remote_storage = HFRemoteStorage( - repo_id="allenai/molmospaces", - repo_prefix=TYPE_TO_PREFIX[args.type], - token=args.hf_token or os.getenv("HF_TOKEN"), - ) + repo_id="allenai/molmospaces", + repo_prefix=TYPE_TO_PREFIX[args.type], + token=args.hf_token or os.getenv("HF_TOKEN"), + ) elif args.storage == "r2": remote_storage = R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") else: diff --git a/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py b/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py index d72e2a98..ce989d9e 100644 --- a/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py +++ b/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py @@ -114,10 +114,10 @@ def main() -> int: if args.storage == "hf": remote_storage = HFRemoteStorage( - repo_id="allenai/molmospaces", - repo_prefix=TYPE_TO_PREFIX[args.type], - token=args.hf_token or os.getenv("HF_TOKEN"), - ) + repo_id="allenai/molmospaces", + repo_prefix=TYPE_TO_PREFIX[args.type], + token=args.hf_token or os.getenv("HF_TOKEN"), + ) elif args.storage == "r2": remote_storage = R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") else: From 01452df8940b5084a7dc8a655b50fbf924f09b0b Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 11 Jun 2026 15:39:15 +0200 Subject: [PATCH 06/11] use r2 by default in molmo_spaces_constants --- molmo_spaces/molmo_spaces_constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index ec3b09b2..72370e73 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -85,7 +85,7 @@ def resource_manager_log_level(log_level=logging.DEBUG): ) STORAGE_IN_USE: Literal["gc", "hf", "r2"] = ( - "gc" # If using hf, HF_TOKEN needs to exist in the environment + "r2" # If using hf, HF_TOKEN needs to exist in the environment ) DATA_TYPE_TO_SOURCE_TO_VERSION = dict( From 038e26afefa269303c62b4e52f8c0132bf86305a Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 18 Jun 2026 10:56:45 +0200 Subject: [PATCH 07/11] Use google cloud storage --- molmo_spaces/molmo_spaces_constants.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index 72370e73..6ec1e6f4 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -84,9 +84,8 @@ def resource_manager_log_level(log_level=logging.DEBUG): else None ) -STORAGE_IN_USE: Literal["gc", "hf", "r2"] = ( - "r2" # If using hf, HF_TOKEN needs to exist in the environment -) +# If using hf, HF_TOKEN needs to exist in the environment +STORAGE_IN_USE: Literal["gc", "hf", "r2"] = "gc" DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ From b87de5404a6f9a33e2d07520ec24b94df95a144b Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 18 Jun 2026 14:19:37 +0200 Subject: [PATCH 08/11] storage type from environment --- molmo_spaces/molmo_spaces_constants.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index 6ec1e6f4..dd047556 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -84,8 +84,18 @@ def resource_manager_log_level(log_level=logging.DEBUG): else None ) -# If using hf, HF_TOKEN needs to exist in the environment -STORAGE_IN_USE: Literal["gc", "hf", "r2"] = "gc" +# If using hf, HF_TOKEN may need to exist in the environment +_valid_storage_types = {"gc", "hf", "r2"} + +DEFAULT_STORAGE: Literal["gc", "hf", "r2"] = "gc" + +_raw_storage_type = os.getenv("MLSPACES_RESOURCE_STORAGE", DEFAULT_STORAGE) +if _raw_storage_type not in _valid_storage_types: + raise ValueError( + f"Invalid MLSPACES_RESOURCE_STORAGE={_raw_storage_type!r}. Must be one of: {_valid_storage_types!r}" + ) + +STORAGE_IN_USE: Storage = _raw # type: ignore[assignment] DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ From 8659e98ec9a3d7c4e61aac240015d504686092c9 Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 18 Jun 2026 14:21:27 +0200 Subject: [PATCH 09/11] Update molmo_spaces_constants.py --- molmo_spaces/molmo_spaces_constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index dd047556..cc14ac2c 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -95,7 +95,7 @@ def resource_manager_log_level(log_level=logging.DEBUG): f"Invalid MLSPACES_RESOURCE_STORAGE={_raw_storage_type!r}. Must be one of: {_valid_storage_types!r}" ) -STORAGE_IN_USE: Storage = _raw # type: ignore[assignment] +STORAGE_IN_USE: Storage = _raw_storage_type # type: ignore[assignment] DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ From 3c43f7ee801a0f11b1b5bbea7696e974cd1702bd Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 18 Jun 2026 14:24:36 +0200 Subject: [PATCH 10/11] typo --- molmo_spaces/molmo_spaces_constants.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index cc14ac2c..725d8b1e 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -87,7 +87,9 @@ def resource_manager_log_level(log_level=logging.DEBUG): # If using hf, HF_TOKEN may need to exist in the environment _valid_storage_types = {"gc", "hf", "r2"} -DEFAULT_STORAGE: Literal["gc", "hf", "r2"] = "gc" +StorageTypes = Literal["gc", "hf", "r2"] + +DEFAULT_STORAGE: StorageTypes = "gc" _raw_storage_type = os.getenv("MLSPACES_RESOURCE_STORAGE", DEFAULT_STORAGE) if _raw_storage_type not in _valid_storage_types: @@ -95,7 +97,7 @@ def resource_manager_log_level(log_level=logging.DEBUG): f"Invalid MLSPACES_RESOURCE_STORAGE={_raw_storage_type!r}. Must be one of: {_valid_storage_types!r}" ) -STORAGE_IN_USE: Storage = _raw_storage_type # type: ignore[assignment] +STORAGE_IN_USE: StorageTypes = _raw_storage_type # type: ignore[assignment] DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ From d0a70a63bf78f4ae835b15b40018186069e9164c Mon Sep 17 00:00:00 2001 From: Jordi Salvador Date: Thu, 18 Jun 2026 14:35:35 +0200 Subject: [PATCH 11/11] Update README.md --- README.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 0ec66c01..cd77ab32 100644 --- a/README.md +++ b/README.md @@ -127,12 +127,13 @@ pip install -e ".[mujoco,curobo]" You may wish to specify some environment variables to configure behavior. Environment variables beginning with the `MLSPACES` prefix can be used to customize MolmoSpaces behavior. -| Environment Variable | Effect | Default | -|---|---|---| -| `MLSPACES_ASSETS_DIR` | Where to place downloaded assets | `~/.cache/molmospaces/assets/` | -| `MLSPACES_FORCE_INSTALL` | Override existing assets | `True` | -| `MLSPACES_PINNED_ASSETS_FILE` | A `.json` file containing pinned versions for each asset, used to override the versions specified in [molmo_spaces_constants.py](molmo_spaces/molmo_spaces_constants.py). | | -| `MUJOCO_EGL_DEVICE_ID` | The rendering device; indices do not always match `CUDA_VISIBLE_DEVICES`. See [here](https://github.com/allenai/molmospaces/issues/66) for details. | `0`| +| Environment Variable | Effect | Default | +|-------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------| +| `MLSPACES_ASSETS_DIR` | Where to place downloaded assets | `~/.cache/molmospaces/assets/` | +| `MLSPACES_FORCE_INSTALL` | Override existing assets | `True` | +| `MLSPACES_PINNED_ASSETS_FILE` | A `.json` file containing pinned versions for each asset, used to override the versions specified in [molmo_spaces_constants.py](molmo_spaces/molmo_spaces_constants.py). | | +| `MUJOCO_EGL_DEVICE_ID` | The rendering device; indices do not always match `CUDA_VISIBLE_DEVICES`. See [here](https://github.com/allenai/molmospaces/issues/66) for details. | `0` | +| `MLSPACES_RESOURCE_STORAGE` | Where to download assets from. Options: `{gc, hf, r2}` for Google Cloud, HuggingFace, or Cloudflare R2 | `gc` | ### Quick Test