33from pathlib import Path
44from typing import Dict , Optional , Tuple , Union
55
6+ from jaxfluids_rl .jxf_env import JAXFluidsEnv , RenderMode
67from omegaconf import OmegaConf
78
89from hydrogym .data_manager import HFDataManager
910
10- from jaxfluids_rl .jxf_env import JAXFluidsEnv , RenderMode
11-
1211
1312class ConfigError (Exception ):
1413 """Exception raised for configuration-related errors."""
@@ -21,7 +20,7 @@ class JAXFluidsFlowEnv(JAXFluidsEnv):
2120 Base JAXFluidsFlowEnv with Hugging Face Hub integration for configuration management.
2221
2322 Arguments:
24- - environment_name: Required. Name of the enviroment .
23+ - environment_name: Required. Name of the environment .
2524 - hf_repo_id: Hugging Face repository (default: 'dynamicslab/HydroGym-environments')
2625
2726 - use_clean_cache: Use clean cache directory (default: True)
@@ -35,7 +34,6 @@ class JAXFluidsFlowEnv(JAXFluidsEnv):
3534 """
3635
3736 def _init_from_hf (self , env_config : dict ) -> None :
38-
3937 # Initialize HF data manager
4038 self .hf_repo_id = env_config .get ("hf_repo_id" , "dynamicslab/HydroGym-environments" )
4139 self .local_fallback_dir = env_config .get ("local_fallback_dir" , None )
@@ -45,15 +43,15 @@ def _init_from_hf(self, env_config: dict) -> None:
4543 repo_id = self .hf_repo_id ,
4644 local_fallback_dir = self .local_fallback_dir ,
4745 use_clean_cache = self .use_clean_cache ,
48- fallback_profile = "JAXFLUIDS"
46+ fallback_profile = "JAXFLUIDS" ,
4947 )
5048
5149 # Environment identification
5250 self .environment_name = env_config .get ("environment_name" )
5351
5452 if not self .environment_name :
5553 raise ConfigError ("'environment_name' must be specified in env_config" )
56-
54+
5755 # Download/get environment configuration
5856 self .env_data_path = self ._setup_environment_data ()
5957
@@ -65,11 +63,10 @@ def _init_from_hf(self, env_config: dict) -> None:
6563 f"No configuration file found for environment '{ self .environment_name } '. "
6664 f"Expected config.yaml in: { self .env_data_path } "
6765 )
68-
66+
6967 # Load configuration from HF
7068 self .conf = OmegaConf .load (self .configuration_file )
7169
72-
7370 def _setup_environment_data (self ) -> str :
7471 """
7572 Download and setup environment data from HF Hub.
@@ -95,7 +92,6 @@ def _setup_environment_data(self) -> str:
9592 return env_path
9693 except Exception as e :
9794 raise ConfigError (f"Failed to setup environment data for { self .environment_name } : { e } " )
98-
9995
10096 def _resolve_configuration_file (self , config_file_input : Optional [str ]) -> Optional [str ]:
10197 """
@@ -154,7 +150,7 @@ def _resolve_configuration_file(self, config_file_input: Optional[str]) -> Optio
154150 f" - Current directory: { os .getcwd ()} \n "
155151 f" - Environment directory: { self .env_data_path } "
156152 )
157-
153+
158154 def _find_configuration_file (self ) -> Optional [str ]:
159155 """
160156 Auto-detect configuration file in the environment data directory.
@@ -192,4 +188,4 @@ def _find_configuration_file(self) -> Optional[str]:
192188 if os .path .exists (self .env_data_path ):
193189 print (f"Available files: { os .listdir (self .env_data_path )} " )
194190
195- return None
191+ return None
0 commit comments