Skip to main content

Core Jax Gym Environment Module

This module provides the base environment class for CFD reinforcement learning with Hugging Face Hub integration for configuration management.

ConfigError Objects​

class ConfigError(Exception)

Exception raised for configuration-related errors.

JAXFlowEnv Objects​

class JAXFlowEnv(environment.Environment[EnvState, EnvParams])

Base JAXFlowEnv with Hugging Face Hub integration for configuration management.

This environment provides a Gymnax-compatible interface for CFD simulations using JAX solvers. It handles:

  • Environment data management via Hugging Face Hub
  • Configuration file resolution and loading
  • Action space configuration

Attributes:

  • environment_name - Name of the CFD environment configuration.
  • env_data_path - Path to the local environment data directory.
  • cfg - OmegaConf configuration object.
  • observation_space - Gymnax observation space.
  • action_space - Gymnax action space.

__init__​

def __init__(env_config: Dict)

Initialize the JAXFlowEnv environment.

Arguments:

  • env_config - Configuration dictionary containing:
    • environment_name (str): Required. Name of the environment.
    • hf_repo_id (str): HF repository ID. Default: 'dynamicslab/HydroGym-environments'
    • local_fallback_dir (str): Optional local fallback directory.
    • use_clean_cache (bool): Whether to use clean cache. Default: True
    • configuration_file (str): Optional path to config file.
    • is_testing (bool): Whether in testing mode. Default: False
    • probe_locations (list): Probe coordinate locations.
    • obs_normalization_strategy (str): One of 'U_inf', 'probewise_mean_std', 'none', 'customized'.
    • obs_loc (list): Required if strategy is 'customized'.
    • obs_scale (list): Required if strategy is 'customized'.

Raises:

  • ConfigError - If required configuration is missing or invalid.

get_environment_files_info​

def get_environment_files_info() -> Dict

Get information about where environment files are stored.

Useful for debugging and understanding file locations.

Returns:

Dictionary containing environment name, paths, and file information.

create_from_hf_env​

@classmethod
def create_from_hf_env(cls,
                       environment_name: str,
                       hf_repo_id: str = "your-username/maiagym-envs",
                       local_fallback_dir: Optional[str] = None,
                       **kwargs)

Create environment directly from HF environment name.

Arguments:

  • environment_name - Name of environment (e.g., 'Cylinder_2D_Re200').
  • hf_repo_id - Hugging Face repository ID.
  • local_fallback_dir - Local fallback directory.
  • **kwargs - Additional environment configuration parameters.

Returns:

Configured MaiaFlowEnv instance.

get_available_environments​

def get_available_environments() -> List[str]

Get list of all available environments from HF Hub.

Returns:

List of environment names.

update_environment_data​

def update_environment_data(force_download: bool = False) -> None

Update environment data from HF Hub.

Arguments:

  • force_download - Force re-download even if cached.

JAXFlowEnvBase Objects​

class JAXFlowEnvBase(environment.Environment[EnvState, EnvParams])

Base JAXFlowEnv without Hugging Face Hub integration. Contains core environment interface methods required by Gymnax.

GymnaxWrapper Objects​

class GymnaxWrapper(object)

Base class for Gymnax wrappers.

FlattenObservationWrapper Objects​

class FlattenObservationWrapper(GymnaxWrapper)

Flatten the observations of the environment.

LogWrapper Objects​

class LogWrapper(GymnaxWrapper)

Log the episode returns and lengths.

ClipAction Objects​

class ClipAction(GymnaxWrapper)

step​

def step(key, state, action, params=None)

TODO: In theory the below line should be the way to do this.