Skip to main content

Core MaiaGym 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.

MaiaFlowEnv Objects

class MaiaFlowEnv(gym.Env)

Base MaiaFlowEnv with Hugging Face Hub integration for configuration management.

This environment provides a Gymnasium-compatible interface for CFD simulations using the m-AIA solver. It handles:

  • Environment data management via Hugging Face Hub
  • Configuration file resolution and loading
  • MPI communication with the CFD solver
  • Observation normalization strategies
  • 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.
  • maiaInterface - MPI interface for communication with m-AIA.
  • observation_space - Gymnasium observation space.
  • action_space - Gymnasium action space.

__init__

def __init__(env_config: Dict)

Initialize the MaiaFlowEnv 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.

configure_observations

def configure_observations() -> None

Configure the number of observations based on observation type and probes.

Sets self.num_outputs based on which observation types are enabled (forces, u, v, w, rho, p) and the number of probes.

setup_normalization

def setup_normalization() -> None

Setup observation normalization factors based on the selected strategy.

Strategies:

  • 'U_inf': Normalize by freestream velocity
  • 'none': No normalization (loc=0, scale=1)
  • 'probewise_mean_std': Compute mean/std from simulation episodes
  • 'customized': Use user-provided loc and scale values

Raises:

  • ConfigError - If normalization configuration is invalid.

compute_normalization_factors

def compute_normalization_factors(zero_actuation: bool = False) -> None

Compute normalization factors for observations by running simulation episodes.

Runs several episodes with random (or zero) actuation to collect probe data, then computes mean and standard deviation for normalization.

Arguments:

  • zero_actuation - If True, use zero actuation. Otherwise, use random actuation.

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.

step

def step(
action: Optional[np.ndarray] = None
) -> Tuple[np.ndarray, float, bool, bool, Dict]

Advance the state of the environment by one step.

Arguments:

  • action - Action array with values in [-1, 1], scaled by MAX_CONTROL.

Returns:

Tuple of (observation, reward, terminated, truncated, info).

reset

def reset(seed: Optional[int] = None,
options: Optional[Dict] = None) -> Tuple[np.ndarray, Dict]

Reset the environment to initial state.

Arguments:

  • seed - Optional random seed for reproducibility.
  • options - Optional reset options.

Returns:

Tuple of (initial_observation, info).

compute_nondim_coefficients

def compute_nondim_coefficients(forces: np.ndarray,
density: float = 1.0,
referenceVelocity: float = 0.1 / np.sqrt(3),
projectionLength: float = 20.0) -> np.ndarray

Compute non-dimensionalized force coefficients.

Arguments:

  • forces - Force array from the CFD solver.
  • density - Reference density.
  • referenceVelocity - Reference velocity.
  • projectionLength - Reference projection length.

Returns:

Non-dimensional force coefficients.

convert_action

def convert_action(action: np.ndarray) -> List

Convert RL action to CFD actuation format.

This method should be overridden by subclasses to implement environment-specific action conversion.

Arguments:

  • action - Action array from the RL agent.

Returns:

Action sequence for the CFD solver.

get_reward

def get_reward() -> Tuple[float, Dict]

Compute the reward for the current state.

This method should be overridden by subclasses to implement environment-specific reward computation.

Returns:

Tuple of (reward, objective_dict).

check_complete

def check_complete() -> bool

Check if the episode is complete.

Returns:

True if the episode has reached max_episode_steps.

close

def close() -> None

Close the environment and signal maia to finish.

This sends a finish signal to the maia solver to properly terminate the MPMD coupling.

set_observation_action_spaces

def set_observation_action_spaces() -> None

Set up the observation and action spaces for Gymnasium compatibility.

configure_probe_dimensions

def configure_probe_dimensions() -> None

Configure probe dimensions based on simulation dimensionality.

Sets noProbeVars (number of variables per probe) and noProbes (total probes).

register_environment

def register_environment(env_prefix: str, env_class: type) -> None

Register an environment class for automatic detection.

Arguments:

  • env_prefix - Environment prefix (e.g., 'Cylinder').
  • env_class - Environment class to register.

from_hf

def from_hf(environment_name: str,
hf_repo_id: str = "dynamicslab/HydroGym-environments",
**kwargs) -> MaiaFlowEnv

Factory function to create any MaiaGym environment from Hugging Face Hub.

Arguments:

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

Returns:

Configured environment instance.

Raises:

  • ValueError - If environment name format is invalid or type is unknown.
  • ConfigError - If environment creation fails.

Examples:

>>> env = maiaGym.from_hf('Cylinder_2D_Re200') >>> env = maiaGym.from_hf('RotaryCylinder_2D_Re1000', obs_loc=custom_loc) >>> env = maiaGym.from_hf('Cavity_2D_Re4140')

list_available_environments

def list_available_environments(
hf_repo_id: str = "dynamicslab/HydroGym-environments") -> List[str]

List all available environments from HF Hub.

Arguments:

  • hf_repo_id - Hugging Face repository ID.

Returns:

List of environment names.

list_registered_types

def list_registered_types() -> Dict[str, type]

List all registered environment types.

Returns:

Dictionary mapping environment types to their classes.