diff --git a/examples/nek/.gitignore b/examples/nek/.gitignore new file mode 100644 index 00000000..ec3c1de5 --- /dev/null +++ b/examples/nek/.gitignore @@ -0,0 +1,9 @@ +# Runtime log and output files +log.run +log.test +log.train +getting_started/*/train_run/ +*.f[0-9]* +*.npz +core.* +SESSION.NAME diff --git a/examples/nek/getting_started/1_nekenv_single/run_nekenv_docker.sh b/examples/nek/getting_started/1_nekenv_single/run_nekenv_docker.sh index 5023348b..70a194d7 100755 --- a/examples/nek/getting_started/1_nekenv_single/run_nekenv_docker.sh +++ b/examples/nek/getting_started/1_nekenv_single/run_nekenv_docker.sh @@ -42,9 +42,10 @@ if [ "$MODE" == "train" ]; then --work-dir "$WORK_DIR" \ --cache-dir "$HOME/.cache/hydrogym" + cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../train_sb3_nek_direct.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ @@ -63,11 +64,12 @@ else --local-dir "$LOCAL_DIR" \ --env "$ENV_NAME" \ --work-dir "$WORK_DIR" \ - --cache-dir "$HOME/.cache/hydrogym" + --cache-dir "$HOME/.cache/hydrogym" \ + --restart-index 1 cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../test_nek_direct.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ diff --git a/examples/nek/getting_started/1_nekenv_single/test_nek_direct.py b/examples/nek/getting_started/1_nekenv_single/test_nek_direct.py index 4f3712bb..bb9bcf1a 100755 --- a/examples/nek/getting_started/1_nekenv_single/test_nek_direct.py +++ b/examples/nek/getting_started/1_nekenv_single/test_nek_direct.py @@ -25,9 +25,10 @@ - Zero control: action = 0 (baseline test) """ -import sys import argparse +import sys from pathlib import Path + import numpy as np from hydrogym.nek import NekEnv @@ -58,6 +59,13 @@ def main(): # Direct instantiation env = NekEnv(env_config=env_config) + # Modify the par file to ensure the simulation configuration is correct + from hydrogym.nek.nek_lib.nek_utils import NEK_INIT + + nek_init = NEK_INIT(nek=env.conf.simulation, drl=env.conf.runner, rank_folder=env.run_folder) + nek_init.rewrite_REA_v19() # Rewrite the par file, v19 corresponds to the new Nek5000 format + # The simulation will be reset, so the par file is to be written out at this point + print("\nEnvironment info:") print("=" * 80) print(f" Observation space: {env.observation_space.shape}") @@ -75,17 +83,17 @@ def main(): print(f"\nRunning {max_steps} steps with zero control...") total_reward = 0.0 - action_dim = env.action_space.shape[0] + # action_dim = env.action_space.shape[0] --> uncomment when needed for step in range(max_steps): # Define action (example: zero control - baseline) - action = np.zeros(action_dim, dtype=np.float32) + # action = np.zeros(action_dim, dtype=np.float32) # OR use constant blowing: # action = np.ones(action_dim, dtype=np.float32) * 0.01 - # OR use opposition control: - # action = -obs[:action_dim] + # Oppose to the wall-normal vel, as the observation is staggered so we sort the even indices + action = -obs[1::2] # Step environment obs, reward, terminated, truncated, info = env.step(action) diff --git a/examples/nek/getting_started/1_nekenv_single/train_sb3_nek_direct.py b/examples/nek/getting_started/1_nekenv_single/train_sb3_nek_direct.py index 816ad0ba..17c393d4 100755 --- a/examples/nek/getting_started/1_nekenv_single/train_sb3_nek_direct.py +++ b/examples/nek/getting_started/1_nekenv_single/train_sb3_nek_direct.py @@ -4,10 +4,10 @@ Includes Monitor, DummyVecEnv, TensorBoard, and VecNormalize best practices. """ -import sys import argparse -from pathlib import Path +import sys from datetime import datetime +from pathlib import Path from hydrogym.nek import NekEnv @@ -28,9 +28,9 @@ def train_single_agent(args): # Import SB3 components try: + from stable_baselines3.common.callbacks import CheckpointCallback from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize - from stable_baselines3.common.callbacks import CheckpointCallback if args.algo == "PPO": from stable_baselines3 import PPO as Algorithm @@ -53,6 +53,12 @@ def make_env(): "configuration_file": args.config_file, } env = NekEnv(env_config=env_config) + + # Modify the par file to ensure the simulation configuration is correct before training + from hydrogym.nek.nek_lib.nek_utils import NEK_INIT + + nek_init = NEK_INIT(nek=env.conf.simulation, drl=env.conf.runner, rank_folder=env.run_folder) + nek_init.rewrite_REA_v19() # Rewrite the par file, v19 corresponds to the new Nek5000 format env = Monitor(env) # CRITICAL: Enables episode reward/length logging return env @@ -61,6 +67,8 @@ def make_env(): # 3. Apply VecNormalize (Crucial for Fluid Dynamics) # This scales inputs to mean 0, std 1 so the Neural Net learns faster. + # VecNormalize is not used in the literature, so it is not guaranteed to work. + # Please see MARL set for more details. env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0) print("Environment created (Wrapped in Monitor, DummyVecEnv, VecNormalize):") diff --git a/examples/nek/getting_started/2_parallel_env/run_parallel_docker.sh b/examples/nek/getting_started/2_parallel_env/run_parallel_docker.sh index 2c8b834b..6fbd2f6f 100755 --- a/examples/nek/getting_started/2_parallel_env/run_parallel_docker.sh +++ b/examples/nek/getting_started/2_parallel_env/run_parallel_docker.sh @@ -22,7 +22,7 @@ WORK_DIR="./train_run" LOCAL_DIR="/workspace/hydrogym/packaged_envs" ENV_NAME="TCFmini_3D_Re180" NPROC_NEK=10 -NUM_STEPS=100 +NUM_STEPS=1000 TOTAL_TIMESTEPS=50000 MODE="${1:-test}" # test or train @@ -44,7 +44,7 @@ if [ "$MODE" == "train" ]; then cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../train_sb3_parallel.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ @@ -66,7 +66,7 @@ else cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../test_nek_parallel.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ diff --git a/examples/nek/getting_started/2_parallel_env/test_nek_parallel.py b/examples/nek/getting_started/2_parallel_env/test_nek_parallel.py index b68fe71f..360718df 100755 --- a/examples/nek/getting_started/2_parallel_env/test_nek_parallel.py +++ b/examples/nek/getting_started/2_parallel_env/test_nek_parallel.py @@ -23,12 +23,14 @@ - Each agent receives observations from sensors near its actuator """ -import sys import argparse +import sys from pathlib import Path + import numpy as np from hydrogym.nek import NekEnv +from hydrogym.nek.nek_lib.nek_utils import NEK_INIT from hydrogym.nek.parallel_env import NekParallelEnv @@ -56,6 +58,10 @@ def main(): } base_env = NekEnv(env_config=env_config) + nek_init = NEK_INIT(nek=base_env.conf.simulation, drl=base_env.conf.runner, rank_folder=base_env.run_folder) + + # Rewrite the par file, v19 corresponds to the new Nek5000 format + nek_init.rewrite_REA_v19() # Wrap with parallel multi-agent environment env = NekParallelEnv(base_env) @@ -98,8 +104,8 @@ def main(): # Strategy 2: Uniform blowing (uncomment to test) # actions = {agent: np.ones(1, dtype=np.float32) * 0.01 for agent in env.agents} - # Strategy 3: Opposition control per agent (uncomment to test) - # actions = {agent: -obs_dict[agent][:1] for agent in env.agents} + # Strategy 3: Opposition control per agent, opposing the wall-normal velocity in this case + actions = {agent: -obs_dict[agent][1:] for agent in env.agents} # Step environment obs_dict, rewards_dict, terminated_dict, truncated_dict, infos_dict = env.step(actions) diff --git a/examples/nek/getting_started/2_parallel_env/train_sb3_parallel.py b/examples/nek/getting_started/2_parallel_env/train_sb3_parallel.py index 820858a9..2c0e4672 100755 --- a/examples/nek/getting_started/2_parallel_env/train_sb3_parallel.py +++ b/examples/nek/getting_started/2_parallel_env/train_sb3_parallel.py @@ -10,14 +10,16 @@ For production, see chapter 3 (PettingZoo + SuperSuit). """ -import sys import argparse -from pathlib import Path +import sys from datetime import datetime -import numpy as np +from pathlib import Path + import gymnasium as gym +import numpy as np from hydrogym.nek import NekEnv, NekParallelEnv +from hydrogym.nek.nek_lib.nek_utils import NEK_INIT class CentralizedParallelWrapper(gym.Env): @@ -127,6 +129,11 @@ def train_parallel_centralized(args): } base_env = NekEnv(env_config=env_config) + # Rewrite the parameter file to ensure the simulation configuration is correct, and + # complies with the v19 Nek5000 format + nek_init = NEK_INIT(nek=base_env.conf.simulation, drl=base_env.conf.runner, rank_folder=base_env.run_folder) + nek_init.rewrite_REA_v19() + # Wrap with parallel interface (dict-based) print("Wrapping with NekParallelEnv (dict-based)...") parallel_env = NekParallelEnv(base_env) @@ -143,9 +150,9 @@ def train_parallel_centralized(args): # Import SB3 try: + from stable_baselines3.common.callbacks import CheckpointCallback from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize - from stable_baselines3.common.callbacks import CheckpointCallback if args.algo == "PPO": from stable_baselines3 import PPO as Algorithm diff --git a/examples/nek/getting_started/3_pettingzoo/README.md b/examples/nek/getting_started/3_pettingzoo/README.md index fbb34bbe..78f15f91 100644 --- a/examples/nek/getting_started/3_pettingzoo/README.md +++ b/examples/nek/getting_started/3_pettingzoo/README.md @@ -38,6 +38,46 @@ mpirun -np 1 python test_nek_pettingzoo.py --steps 100 : -np 10 nek5000 mpirun -np 1 python train_sb3_pettingzoo.py --env MiniChannel_Re180 --algo PPO --total-timesteps 100000 : -np 10 nek5000 ``` +## Configuration-Driven Tutorial (Recommended for Reproducibility) + +Use a fixed YAML config to lock simulation + runner settings across runs. + +### 1) Prepare a workspace +```bash +python ../prepare_workspace.py \ + --local-dir ../../../packaged_envs \ + --env TCFmini_3D_Re180 \ + --work-dir ./train_run +``` + +### 2) Train with a config file +```bash +cd train_run +mpirun -np 1 python ../train_sb3_pettingzoo.py \ + --env TCFmini_3D_Re180 \ + --nproc 10 \ + --config-file ../configs/pettingzoo_tcfmini_re180.yml \ + --algo TD3 \ + --total-timesteps 5000000 \ + : -np 10 nek5000 +``` + +### 3) Evaluate (PettingZoo rollouts) +```bash +cd train_run +mpirun -np 1 python ../test_nek_pettingzoo.py \ + --env TCFmini_3D_Re180 \ + --nproc 10 \ + --config-file ../configs/pettingzoo_tcfmini_re180.yml \ + --steps 2500 \ + : -np 10 nek5000 +``` + +Notes: +- The config lives in `examples/nek/configs/pettingzoo_tcfmini_re180.yml`. +- Run from the workspace (`train_run`) so `compile_path: "."` resolves to case files. +- Ensure `--nproc` matches `simulation.nproc` in the config. + ## When to Use - **Production SB3 training** on multi-agent environments diff --git a/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh b/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh index 753c49c8..5c5cc5eb 100755 --- a/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh +++ b/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh @@ -1,6 +1,7 @@ #!/usr/bin/env bash # # Run NEK5000 PettingZoo tests with MPMD coupling. +# Config is loaded automatically from HuggingFace (environment_config.yaml). # # Usage: # ./run_pettingzoo_docker.sh # Test only @@ -44,7 +45,7 @@ if [ "$MODE" == "train" ]; then cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../train_sb3_pettingzoo.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ @@ -66,7 +67,7 @@ else cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../test_nek_pettingzoo.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ diff --git a/examples/nek/getting_started/3_pettingzoo/test_nek_pettingzoo.py b/examples/nek/getting_started/3_pettingzoo/test_nek_pettingzoo.py index 4939b5ad..04117c97 100755 --- a/examples/nek/getting_started/3_pettingzoo/test_nek_pettingzoo.py +++ b/examples/nek/getting_started/3_pettingzoo/test_nek_pettingzoo.py @@ -23,9 +23,10 @@ - Actions and observations are dicts with agent names as keys """ -import sys import argparse +import sys from pathlib import Path + import numpy as np from hydrogym.nek import NekEnv @@ -115,6 +116,8 @@ def main(): # Strategy 3: Cooperative strategy - all agents use same signal # signal = np.sin(step * 0.1) * 0.01 # actions = {agent: np.array([signal], dtype=np.float32) for agent in env.agents} + # Opposition Control Strategy, oppose to the wall-normal velocity (-1) + actions = {agent: -1.0 * obs_dict[agent][:-1] for agent in env.agents} # Step environment obs_dict, rewards_dict, terminated_dict, truncated_dict, infos_dict = env.step(actions) diff --git a/examples/nek/getting_started/3_pettingzoo/train_sb3_pettingzoo.py b/examples/nek/getting_started/3_pettingzoo/train_sb3_pettingzoo.py index 4e5e12e0..949fb8d4 100755 --- a/examples/nek/getting_started/3_pettingzoo/train_sb3_pettingzoo.py +++ b/examples/nek/getting_started/3_pettingzoo/train_sb3_pettingzoo.py @@ -9,10 +9,10 @@ Educational approach (DIY wrapper): See chapter 2 """ -import sys import argparse -from pathlib import Path +import sys from datetime import datetime +from pathlib import Path from hydrogym.nek import NekEnv, make_pettingzoo_env @@ -35,6 +35,10 @@ def train_pettingzoo_with_supersuit(args): "use_clean_cache": False, "local_fallback_dir": args.local_dir, "configuration_file": args.config_file, + "rescale_actions": True, # The action space has range of [-1,1] and we rescale by the maximum amplitude + "normalize_input": "utau", + "ctrl_min_amp": -0.06388353, + "ctrl_max_amp": 0.06388353, } base_env = NekEnv(env_config=env_config) @@ -47,8 +51,8 @@ def train_pettingzoo_with_supersuit(args): # Convert to SB3-compatible format using SuperSuit try: - from supersuit import black_death_v3, pad_observations_v0, pad_action_space_v0 from pettingzoo.utils import parallel_to_aec + from supersuit import black_death_v3, pad_action_space_v0, pad_observations_v0 except ImportError: print("✗ Error: PettingZoo/SuperSuit not installed!") print("Install with: pip install pettingzoo supersuit") @@ -56,9 +60,10 @@ def train_pettingzoo_with_supersuit(args): try: import numpy as np - from stable_baselines3.common.monitor import Monitor - from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize, VecEnvWrapper from stable_baselines3.common.callbacks import CheckpointCallback + from stable_baselines3.common.monitor import Monitor + from stable_baselines3.common.noise import NormalActionNoise + from stable_baselines3.common.vec_env import DummyVecEnv, VecEnvWrapper, VecNormalize if args.algo == "PPO": from stable_baselines3 import PPO as Algorithm @@ -113,7 +118,7 @@ def step_wait(self): # Wrap with VecNormalize print("Wrapping with VecNormalize...") - env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0) + env = VecNormalize(env, norm_obs=False, norm_reward=False, clip_obs=0.0) print("Final environment:") print(f" Observation space: {env.observation_space}") @@ -131,20 +136,33 @@ def step_wait(self): "policy": "MlpPolicy", "env": env, "learning_rate": args.learning_rate, + "policy_kwargs": { + "net_arch": { + "qf": [16, 16, 64], + "pi": [ + 8, + ], + } + }, "gamma": args.gamma, "verbose": 1, "tensorboard_log": str(log_dir), } + # NOTE: PPO is not used in the literature, so it is not guaranteed to work. if args.algo == "PPO": model_kwargs["n_steps"] = args.n_steps model_kwargs["batch_size"] = args.batch_size else: model_kwargs["batch_size"] = args.batch_size - model_kwargs["buffer_size"] = 100000 + model_kwargs["buffer_size"] = int(1e6) model_kwargs["learning_starts"] = 100 - model_kwargs["train_freq"] = 1 - model_kwargs["gradient_steps"] = 1 + model_kwargs["train_freq"] = (300, "step") + model_kwargs["gradient_steps"] = 64 + model_kwargs["action_noise"] = NormalActionNoise( + mean=np.zeros(env.action_space.shape[0]), # Zero-Mean + sigma=0.1 * np.ones(env.action_space.shape[0]), # Action noise is 10% of the action space + ) model = Algorithm(**model_kwargs) print("✓ Model created\n") @@ -152,6 +170,9 @@ def step_wait(self): # Calculate safe save frequency safe_save_freq = max(args.save_freq // env.num_envs, 1) + # CallBack list during training process + callback_list = [] + # Custom callback to save VecNormalize stats class SaveVecNormalizeCallback(CheckpointCallback): def _on_step(self) -> bool: @@ -165,6 +186,11 @@ def _on_step(self) -> bool: save_freq=safe_save_freq, save_path=str(log_dir), name_prefix="model" ) + callback_list.append(checkpoint_callback) + + # Add CheckpointCallback to the callback list + callback_list.append(CheckpointCallback(save_freq=safe_save_freq, save_path=str(log_dir), name_prefix="rl-model")) + print(f"Log directory: {log_dir}\n") # Train @@ -173,7 +199,7 @@ def _on_step(self) -> bool: print("=" * 70 + "\n") try: - model.learn(total_timesteps=args.total_timesteps, callback=checkpoint_callback, tb_log_name=f"{args.algo}_run") + model.learn(total_timesteps=args.total_timesteps, callback=callback_list, tb_log_name=f"{args.algo}_run") # Save final model and normalization stats final_model_path = log_dir / "model_final.zip" @@ -208,13 +234,13 @@ def main(): parser.add_argument("--algo", default="PPO", choices=["PPO", "TD3", "SAC"]) parser.add_argument("--total-timesteps", type=int, default=100000) parser.add_argument("--n-steps", type=int, default=2048) - parser.add_argument("--learning-rate", type=float, default=3e-4) - parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--learning-rate", type=float, default=1e-3) + parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--gamma", type=float, default=0.99) # Logging parser.add_argument("--log-dir", default="./logs") - parser.add_argument("--save-freq", type=int, default=10000) + parser.add_argument("--save-freq", type=int, default=5) args = parser.parse_args() diff --git a/examples/nek/getting_started/4_from_hf/run_from_hf_docker.sh b/examples/nek/getting_started/4_from_hf/run_from_hf_docker.sh index 5ee83664..7ed29a3f 100755 --- a/examples/nek/getting_started/4_from_hf/run_from_hf_docker.sh +++ b/examples/nek/getting_started/4_from_hf/run_from_hf_docker.sh @@ -44,7 +44,7 @@ if [ "$MODE" == "train" ]; then cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../train_sb3_from_hf.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ @@ -66,7 +66,7 @@ else cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../test_nek_DM.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ diff --git a/examples/nek/getting_started/4_from_hf/test_nek_DM.py b/examples/nek/getting_started/4_from_hf/test_nek_DM.py index 0c270540..30e7c5f4 100755 --- a/examples/nek/getting_started/4_from_hf/test_nek_DM.py +++ b/examples/nek/getting_started/4_from_hf/test_nek_DM.py @@ -21,12 +21,14 @@ - Opposition control: action = -observation """ -import sys import argparse +import sys from pathlib import Path + import numpy as np from hydrogym.nek import NekEnv +from hydrogym.nek.nek_lib.nek_utils import NEK_INIT def main(): @@ -46,6 +48,10 @@ def main(): local_fallback_dir=args.local_dir, ) + # Rewrite the par file to ensure the simulation configuration is correct + nek_init = NEK_INIT(nek=env.conf.simulation, drl=env.conf.runner, rank_folder=env.run_folder) + nek_init.rewrite_REA_v19() + print("\nEnvironment info:") print("=" * 80) print(f" Observation space: {env.observation_space.shape}") diff --git a/examples/nek/getting_started/4_from_hf/train_sb3_from_hf.py b/examples/nek/getting_started/4_from_hf/train_sb3_from_hf.py index ed9b3580..263f5cca 100755 --- a/examples/nek/getting_started/4_from_hf/train_sb3_from_hf.py +++ b/examples/nek/getting_started/4_from_hf/train_sb3_from_hf.py @@ -8,12 +8,13 @@ Pattern: NekEnv.from_hf(env_name, nproc, ...) """ -import sys import argparse -from pathlib import Path +import sys from datetime import datetime +from pathlib import Path from hydrogym.nek import NekEnv +from hydrogym.nek.nek_lib.nek_utils import NEK_INIT def train_with_from_hf(args): @@ -35,6 +36,10 @@ def train_with_from_hf(args): local_fallback_dir=args.local_dir, ) + # Rewrite the par file to ensure the simulation configuration is correct + nek_init = NEK_INIT(nek=env.conf.simulation, drl=env.conf.runner, rank_folder=env.run_folder) + nek_init.rewrite_REA_v19() + print("\nEnvironment created:") print(f" Observation space: {env.observation_space.shape}") print(f" Action space: {env.action_space.shape}") @@ -42,11 +47,12 @@ def train_with_from_hf(args): # Import SB3 try: + from stable_baselines3.common.callbacks import CheckpointCallback from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize - from stable_baselines3.common.callbacks import CheckpointCallback if args.algo == "PPO": + # NOTE: PPO is not used in the literature, so it is not guaranteed to work. from stable_baselines3 import PPO as Algorithm elif args.algo == "TD3": from stable_baselines3 import TD3 as Algorithm @@ -60,6 +66,7 @@ def train_with_from_hf(args): print("Wrapping with Monitor, DummyVecEnv, VecNormalize...") env = Monitor(env) env = DummyVecEnv([lambda: env]) + # NOTE: Add VecNormalize is not used in the literature, so it is not guaranteed to work. env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0) print("Environment wrapped:") diff --git a/examples/nek/getting_started/5_hydrogym_control/run_control_docker.sh b/examples/nek/getting_started/5_hydrogym_control/run_control_docker.sh index 40acbaa9..fffef3d4 100755 --- a/examples/nek/getting_started/5_hydrogym_control/run_control_docker.sh +++ b/examples/nek/getting_started/5_hydrogym_control/run_control_docker.sh @@ -53,7 +53,7 @@ if [ "$MODE" == "train" ]; then cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../train_sb3_with_integrate.py \ --env "$ENV_NAME" \ --local-dir "$LOCAL_DIR" \ @@ -77,7 +77,7 @@ else cd "$WORK_DIR" || exit 1 - mpirun \ + mpirun --use-hwthread-cpus\ -np 1 python ../test_nek_env_controller.py \ --env "$ENV_NAME" \ --nproc ${NPROC_NEK} \ diff --git a/examples/nek/getting_started/5_hydrogym_control/test_nek_env_controller.py b/examples/nek/getting_started/5_hydrogym_control/test_nek_env_controller.py index 822d1655..34046bac 100755 --- a/examples/nek/getting_started/5_hydrogym_control/test_nek_env_controller.py +++ b/examples/nek/getting_started/5_hydrogym_control/test_nek_env_controller.py @@ -23,20 +23,21 @@ --local-dir: Local fallback directory for environments """ -import sys import argparse import logging +import sys from pathlib import Path from typing import Optional import numpy as np +from hydrogym.nek import NekEnv +from hydrogym.nek.nek_lib.nek_utils import NEK_INIT + # Force unbuffered output for MPMD mode sys.stdout = open(sys.stdout.fileno(), "w", buffering=1) sys.stderr = open(sys.stderr.fileno(), "w", buffering=1) -from hydrogym.nek import NekEnv # noqa: E402 - def setup_logging(verbose: bool = False) -> logging.Logger: """Configure logging with timestamps.""" @@ -56,7 +57,7 @@ def create_controller(env, controller_type: str, logger: logging.Logger): def controller(t, obs, env): # Oppose the velocity at sensor location (obs is flattened 1D array) # For opposition control, use negative of observation values - return -obs[: env.action_space.shape[0]] + return -obs[1::2] # YW: Again note that we are opposing the wall-normal velocity (index = 1) elif controller_type.upper() == "BL": # Constant blowing def controller(t, obs, env): @@ -119,6 +120,7 @@ def run_nek_test( # Legacy pattern with config file logger.info(f"Using legacy pattern with config file: {config_path}") from omegaconf import OmegaConf + from hydrogym.nek.configs import Config # Load config @@ -139,6 +141,10 @@ def run_nek_test( else: raise ValueError("Must provide either --env (MAIA pattern) or --config (legacy pattern)") + # Rewrite the par file to ensure the simulation configuration is correct + nek_init = NEK_INIT(nek=env.conf.simulation, drl=env.conf.runner, rank_folder=env.run_folder) + nek_init.rewrite_REA_v19() + logger.info("✓ Environment created successfully") except Exception as e: logger.error(f"✗ Failed to create environment: {e}") diff --git a/examples/nek/getting_started/5_hydrogym_control/train_sb3_with_integrate.py b/examples/nek/getting_started/5_hydrogym_control/train_sb3_with_integrate.py index 7ea55a05..eca50c6c 100755 --- a/examples/nek/getting_started/5_hydrogym_control/train_sb3_with_integrate.py +++ b/examples/nek/getting_started/5_hydrogym_control/train_sb3_with_integrate.py @@ -18,13 +18,14 @@ mpirun -np 1 python train_sb3_with_integrate.py --config config.yml --nproc 10 : -np 10 nek5000 """ -import sys import argparse -from pathlib import Path +import sys from datetime import datetime +from pathlib import Path from typing import Optional import numpy as np + from hydrogym.nek import NekEnv, integrate @@ -62,6 +63,7 @@ def create_environment( # Legacy pattern with config file print(f" Using legacy pattern with config: {config_path}") from omegaconf import OmegaConf + from hydrogym.nek.configs import Config config = OmegaConf.merge( diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/README.md b/examples/nek/getting_started/6_zeroshot_wing_demo/README.md new file mode 100644 index 00000000..77d9df91 --- /dev/null +++ b/examples/nek/getting_started/6_zeroshot_wing_demo/README.md @@ -0,0 +1,117 @@ +# Zero-Shot Wing Deployment (Multi-Policy MARL) + +Zero-shot deployment demo for the small NACA4412 wing case: multiple control policies are mapped to actuator subsets and executed together in one PettingZoo rollout. + +__This is a deployment/evaluation demo only (no training). The template and controllers are intended for demonstration and should not be used to draw physical conclusions.__ + +> NOTE: The provided Nek5000 executable is pre-compiled for this chapter, so this demo focuses on the DRL-style rollout/deployment workflow. + +## What the script does + +`test_nek_pettingzoo.py`: +- loads a base `NekEnv` via `NekEnv.from_hf(...)` and wraps it with `make_pettingzoo_env(...)` +- builds one controller per entry in `POLICY_SPECS` (from `meta_policy_small_wing_template.py`) +- assigns each controller to actuator agents by `x_range` and `side` (`SS` means `y > 0`, `PS` means `y < 0`) +- refreshes each group's actions every `drl_step` steps (refresh at step `0`; otherwise actions are held) +- clips actions to `action_bounds` +- computes an “inverted” + scaled reward summary for display (deployment-only) + +Unassigned actuator agents receive zero action. + +## Interface (PettingZoo rollout) + +```python +from hydrogym.nek import NekEnv +from hydrogym.nek.pettingzoo_env import make_pettingzoo_env + +base_env = NekEnv.from_hf("NACA4412_3D_Re75000_AOA5", nproc=12) +env = make_pettingzoo_env(base_env) + +obs_dict, info = env.reset() +actions = {agent: controller(obs_dict[agent]) for agent in env.agents} +obs_dict, rewards_dict, terminations, truncations, infos = env.step(actions) +``` + +## Files + +- `test_nek_pettingzoo.py` - zero-shot multi-policy rollout demo (deployment only) +- `meta_policy_small_wing_template.py` - template defining `ENV_NAME`, `NPROC`, and `POLICY_SPECS` +- `run_pettingzoo_docker.sh` - runner script (module load + workspace prep + `mpirun`) + +## Usage + +### Recommended: use the runner script +From `6_zeroshot_wing_demo/`: + +```bash +./run_pettingzoo_docker.sh +./run_pettingzoo_docker.sh --policy-root /workspace/legacy_runs +``` + +### Direct: run the Python deployment script + +Default template: +```bash +mpirun -np 1 python test_nek_pettingzoo.py : -np 12 nek5000 +``` + +Legacy policy template + run root: +```bash +mpirun -np 1 python test_nek_pettingzoo.py \ + --policy-template ./meta_policy_small_wing_template.py \ + --policy-root /path/to/legacy_runs \ + --steps 3000 \ + : -np 12 nek5000 +``` + +Useful overrides: +- `--policy-template PATH` (defaults to `./meta_policy_small_wing_template.py`) +- `--env ENV_NAME` (defaults from template `ENV_NAME`) +- `--nproc NPROC` (defaults from template `NPROC`) +- `--steps NUM_STEPS` (defaults from template `NUM_STEPS`) +- `--policy-root PATH` (where RL model run folders live) +- `--local-dir PATH` (optional fallback dir for packaged envs) +- `--log-every N` (reward table frequency) + +## Policy Template (`meta_policy_small_wing_template.py`) + +The template defines a lightweight legacy-`MetaPolicy.py`-style configuration. + +Required top-level variables: +- `ENV_NAME` +- `NPROC` +- `NUM_STEPS` +- `POLICY_ROOT` (default for `--policy-root`) +- `POLICY_SPECS` (list of policy group dicts) + +Each `POLICY_SPECS` entry supports: +- `name` +- `x_range: [x_min, x_max]` +- `side: "SS"` (y>0) or `"PS"` (y<0) +- `algorithm: "PPO" | "TD3" | "DDPG" | "BL" | "ZERO"` +- `drl_step` (action refresh interval; actions are held between refreshes) +- `action_bounds: [min, max]` +- optional scaling knobs: `u_tau`, `baseline_dudy` +- RL algorithms only: `agent_run_name`, `policy`, and/or `model_path` + +Algorithm semantics: +- `ZERO` outputs an all-zero action (no model needed) +- `BL` outputs a constant action equal to `action_max` (no model needed) +- `PPO`/`TD3`/`DDPG` load a Stable-Baselines3 model from `model_path`/`POLICY_ROOT` + +For overlapping actuator regions, the last-assigned policy takes precedence. + +## Default RL Model Path Convention + +For RL policies (`PPO`, `TD3`, `DDPG`), if `model_path` is not set, the default expected path is: + +```text +//logs/- +``` + +## Notes + +- Deployment-only (evaluation). No training happens in this chapter. +- `drl_step` controls when the controller is queried; between refreshes, the last action is held for the whole group. +- `u_tau` is used to normalize observations before calling the controller (the code comments note that solver-side normalization by `u_tau` should be kept consistent with how the policies were trained). +- This demo uses deterministic controller calls (`controller.predict(..., deterministic=True)`), and displays a reward summary to help compare controller configurations. diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/meta_policy_small_wing_template.py b/examples/nek/getting_started/6_zeroshot_wing_demo/meta_policy_small_wing_template.py new file mode 100644 index 00000000..bcd285da --- /dev/null +++ b/examples/nek/getting_started/6_zeroshot_wing_demo/meta_policy_small_wing_template.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +""" +Template configuration for zero-shot multi-policy deployment on the small wing. + +The structure mirrors the old MetaPolicy-style configuration, but is lightweight +and directly consumable by ``test_nek_pettingzoo.py`` in this folder. +""" + +# Default environment used in this chapter +ENV_NAME = "NACA4412_3D_Re75000_AOA5" + +# Number of Nek5000 worker processes for this case +NPROC = 12 + +# Total rollout steps for demonstration +NUM_STEPS = 3000 + +# Root directory containing legacy policy run folders. +# Example expected layout for RL entries: +# //logs/- +POLICY_ROOT = "./legacy_runs" + +# Per-policy deployment specification. +# Required fields: +# - name: policy label +# - x_range: [x_min, x_max] +# - side: "SS" (y>0) or "PS" (y<0) +# - algorithm: "PPO"/"TD3"/"DDPG" or baseline "BL"/"OC"/"ZERO" +# - drl_step: action refresh interval (action is held between updates) +# - action_bounds: [min, max] +# Optional fields (for RL algorithms): +# - agent_run_name +# - policy +# Optional scaling fields: +# - u_tau +# - baseline_dudy +POLICY_SPECS = [ + { + "name": "CTRL000", + "x_range": [0.23, 0.40], + "side": "SS", + "algorithm": "BL", # This is a constant action policy, no model needed + "agent_run_name": "601031", # This is the agent run name, it is the name of the folder that contains the model + "policy": "rl_model_749700000_steps", # This is the policy name, it is the name of the policy + "drl_step": 4, # This is the action refresh interval, it is the number of steps between each action refresh + "u_tau": 1.0, # This is the u_tau, it is the u_tau of the environment, it is used to normalize the observation + "baseline_dudy": 1135.83, # This is the baseline_dudy, it is the baseline_dudy of the environment + "action_bounds": [-1.0, 1.0], # This is the action bounds, it is the action bounds of the environment + }, + { + "name": "CTRL001", + "x_range": [0.40, 0.50], + "side": "SS", + "algorithm": "BL", + "agent_run_name": "602031", + "policy": "rl_model_716625000_steps", + "drl_step": 5, + "u_tau": 1.0, + "baseline_dudy": 940.90, + "action_bounds": [-1.0, 1.0], + }, + { + "name": "CTRL002", + "x_range": [0.50, 0.70], + "side": "SS", + "algorithm": "BL", + "agent_run_name": "603031", + "policy": "rl_model_217800000_steps", + "drl_step": 6, + "u_tau": 1.0, + "baseline_dudy": 759.46, + "action_bounds": [-1.0, 1.0], + }, + { + "name": "CTRL003", + "x_range": [0.70, 0.88], + "side": "SS", + "algorithm": "BL", + "agent_run_name": "604031", + "policy": "rl_model_2914890000_steps", + "drl_step": 9, + "u_tau": 1.0, + "baseline_dudy": 485.32, + "action_bounds": [-1.0, 1.0], + }, + { + "name": "CTRL004", + "x_range": [0.23, 0.88], + "side": "PS", + "algorithm": "BL", + "agent_run_name": "601031", + "policy": "rl_model_749700000_steps", + "drl_step": 4, + "u_tau": 1.0, + "baseline_dudy": 417.136, + "action_bounds": [-1.0, 1.0], + }, +] diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/run_pettingzoo_docker.sh b/examples/nek/getting_started/6_zeroshot_wing_demo/run_pettingzoo_docker.sh new file mode 100644 index 00000000..6463cad4 --- /dev/null +++ b/examples/nek/getting_started/6_zeroshot_wing_demo/run_pettingzoo_docker.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# +# Run zero-shot wing deployment demo with MPMD coupling. +# +# Usage: +# ./run_pettingzoo_docker.sh +# ./run_pettingzoo_docker.sh --policy-root /workspace/legacy_runs + +set -e + +# Load Nek5000 module +module purge +module load Nek5000/1.0-gompi-2024a-SystemCUDA-SmallWing + +# Activate Python environment +source ~/venvs/hydrogym_cpu/bin/activate + +export OMP_NUM_THREADS=1 + +# Configuration +WORK_DIR="./train_run" +LOCAL_DIR="/workspace/hydrogym/packaged_envs" +ENV_NAME="NACA4412_3D_Re75000_AOA5" +NPROC_NEK=12 +NUM_STEPS=30 +POLICY_TEMPLATE="../meta_policy_small_wing_template.py" +POLICY_ROOT="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --policy-root) + POLICY_ROOT="$2" + shift 2 + ;; + *) + echo "Unknown argument: $1" + exit 1 + ;; + esac +done + +echo "=== Zero-Shot Wing Deployment (PettingZoo) ===" +echo "Environment: $ENV_NAME" +echo "Nek5000 procs: $NPROC_NEK" +echo "Rollout steps: $NUM_STEPS" +if [ -n "$POLICY_ROOT" ]; then + echo "Policy root: $POLICY_ROOT" +else + echo "Policy root: