Chore: pre-commit again.

og_marl/vault_utils/combine_vaults.py

@@ -101,7 +101,7 @@ def combine_vaults(rel_dir: str, vault_name: str, vault_uids: Optional[list[str]
     # check that a subsampled vault by the same name does not already exist
     if check_directory_exists_and_not_empty(f"./{rel_dir}/{new_vault_name}"):
         print(
-            f"Vault '{rel_dir}/{new_vault_name.removesuffix('.vlt')}' already exists. To combine from scratch, please remove the current combined vault from its directory." #noqa
+            f"Vault '{rel_dir}/{new_vault_name.removesuffix('.vlt')}' already exists. To combine from scratch, please remove the current combined vault from its directory."  # noqa
         )
         return new_vault_name
 

og_marl/vault_utils/download_vault.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
+from typing import List, Dict
 
 import os
 import sys
@@ -139,7 +139,7 @@
 }
 
 
-def print_download_options() -> None:
+def print_download_options() -> Dict[str, Dict]:
     print("VAULT_INFO:")
     for source in VAULT_INFO.keys():
         print(f"\t {source}")
@@ -163,7 +163,7 @@ def download_and_unzip_vault(
         f"{dataset_base_dir}/{dataset_source}/{env_name}/{scenario_name}.vlt"
     ):
         print(
-            f"Vault '{dataset_base_dir}/{dataset_source}/{env_name}/{scenario_name}' already exists."
+            f"Vault '{dataset_base_dir}/{dataset_source}/{env_name}/{scenario_name}' already exists."  # noqa
         )
         return f"{dataset_base_dir}/{dataset_source}/{env_name}/{scenario_name}.vlt"
 
@@ -178,7 +178,7 @@ def download_and_unzip_vault(
         print(
             "Dataset from "
             + str(dataset_download_url)
-            + " could not be downloaded. Try entering a different URL, or removing the part which auto-downloads."
+            + " could not be downloaded. Try entering a different URL, or removing the part which auto-downloads."  # noqa
         )
         return f"{dataset_base_dir}/{dataset_source}/{env_name}/{scenario_name}.vlt"
 

og_marl/vault_utils/subsample_similar.py

@@ -101,7 +101,7 @@ def subsample_similar(
     # check that a subsampled vault by the same name does not already exist
     if check_directory_exists_and_not_empty(f"./{new_rel_dir}/{new_vault_name}"):
         print(
-            f"Vault '{new_rel_dir}/{new_vault_name.removesuffix('.vlt')}' already exists. To subsample from scratch, please remove the current subsampled vault from its directory." #noqa
+            f"Vault '{new_rel_dir}/{new_vault_name.removesuffix('.vlt')}' already exists. To subsample from scratch, please remove the current subsampled vault from its directory."  # noqa
         )
         return
 

og_marl/vault_utils/subsample_smaller.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from git import Optional
 
 import jax
 import pickle
@@ -107,21 +107,22 @@ def stitch_vault_from_sampled_episodes_(
 
     for start, end in zip(len_start_end_sample[:, 1], len_start_end_sample[:, 2]):
         sample_experience = jax.tree_util.tree_map(
-            lambda x: x[:, int(start) : int(end + 1), ...], experience
+            lambda x: x[:, int(start) : int(end + 1), ...],  # noqa
+            experience,
         )
         dest_state = buffer_add(dest_state, sample_experience)
 
     timesteps_written = dest_vault.write(dest_state)
 
     print(timesteps_written)
 
-    return timesteps_written
+    return int(timesteps_written)
 
 
 def subsample_smaller_vault(
     vaults_dir: str,
     vault_name: str,
-    vault_uids: list = None,
+    vault_uids: Optional[list] = None,
     target_number_of_transitions: int = 500000,
 ) -> str:
     # check that the vault to be subsampled exists
@@ -141,7 +142,7 @@ def subsample_smaller_vault(
     # check that a subsampled vault by the same name does not already exist
     if check_directory_exists_and_not_empty(f"./{vaults_dir}/{new_vault_name}"):
         print(
-            f"Vault '{vaults_dir}/{new_vault_name.removesuffix('.vlt')}' already exists. To subsample from scratch, please remove the current subsampled vault from its directory." #noqa
+            f"Vault '{vaults_dir}/{new_vault_name.removesuffix('.vlt')}' already exists. To subsample from scratch, please remove the current subsampled vault from its directory."  # noqa
         )
         return f"./{vaults_dir}/{vault_name}"