317 lines
10 KiB
Python
317 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Train PPO agent for CFD flow control.
|
|
|
|
This script trains a Proximal Policy Optimization (PPO) agent to control
|
|
flow around a cylinder using the CFD environment.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import pickle
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
# Set threading layers
|
|
os.environ['MKL_THREADING_LAYER'] = 'GNU'
|
|
os.environ["OMP_NUM_THREADS"] = "8"
|
|
os.environ["MKL_NUM_THREADS"] = "8"
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.nn import Module
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.callbacks import BaseCallback
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
# Add src to path for imports
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
|
|
|
|
from environments import CFDFlowControlEnv
|
|
from config import (
|
|
load_celeris_configs,
|
|
get_model_path,
|
|
get_tensorboard_logdir,
|
|
get_output_path,
|
|
)
|
|
|
|
|
|
class SinActivation(Module):
|
|
"""Sine activation function for neural networks."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return torch.sin(x)
|
|
|
|
|
|
class TensorboardCallback(BaseCallback):
|
|
"""
|
|
Custom callback for logging additional metrics to TensorBoard.
|
|
"""
|
|
|
|
def __init__(self, check_freq: int = 360, verbose: int = 0):
|
|
super().__init__(verbose)
|
|
self.check_freq = check_freq
|
|
self.episode_rewards = []
|
|
self.episode_cd = []
|
|
self.episode_cl = []
|
|
|
|
def _on_step(self) -> bool:
|
|
if self.n_calls % self.check_freq == 0:
|
|
# Extract episode info
|
|
if len(self.locals.get('infos', [])) > 0:
|
|
info = self.locals['infos'][0]
|
|
if 'cd' in info:
|
|
self.logger.record('flow/cd', info['cd'])
|
|
if 'cl' in info:
|
|
self.logger.record('flow/cl', info['cl'])
|
|
if 'reward_cd' in info:
|
|
self.logger.record('reward/cd', info['reward_cd'])
|
|
if 'reward_cl' in info:
|
|
self.logger.record('reward/cl', info['reward_cl'])
|
|
if 'reward_sim' in info:
|
|
self.logger.record('reward/sim', info['reward_sim'])
|
|
return True
|
|
|
|
|
|
def parse_args():
|
|
"""Parse command line arguments."""
|
|
parser = argparse.ArgumentParser(description='Train PPO for CFD control')
|
|
|
|
# Environment settings
|
|
parser.add_argument('--device-id', type=int, default=0,
|
|
help='CUDA device ID for simulation')
|
|
|
|
# Training hyperparameters
|
|
parser.add_argument('--total-timesteps', type=int, default=100,
|
|
help='Number of training iterations (each = n_steps)')
|
|
parser.add_argument('--n-steps', type=int, default=3600,
|
|
help='Steps to collect per training iteration')
|
|
parser.add_argument('--batch-size', type=int, default=360,
|
|
help='Batch size for PPO updates')
|
|
parser.add_argument('--learning-rate', type=float, default=3e-4,
|
|
help='Learning rate')
|
|
parser.add_argument('--gamma', type=float, default=0.99,
|
|
help='Discount factor')
|
|
|
|
# Model settings
|
|
parser.add_argument('--activation', choices=['tanh', 'relu', 'sin'], default='sin',
|
|
help='Activation function for policy network')
|
|
parser.add_argument('--cuda-device', type=int, default=0,
|
|
help='CUDA device for PyTorch training')
|
|
|
|
# Experiment settings
|
|
parser.add_argument('--run-name', type=str, default='ppo_cfd_control',
|
|
help='Name for this training run')
|
|
parser.add_argument('--save-freq', type=int, default=10,
|
|
help='Save model every N iterations')
|
|
parser.add_argument('--eval-episodes', type=int, default=1,
|
|
help='Number of episodes to evaluate')
|
|
|
|
# Resume training
|
|
parser.add_argument('--resume', type=str, default=None,
|
|
help='Path to model checkpoint to resume from')
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def create_env(device_id: int):
|
|
"""Create the CFD environment."""
|
|
config_cuda, config_field = load_celeris_configs()
|
|
|
|
env = CFDFlowControlEnv(
|
|
device_id=device_id,
|
|
config_cuda=config_cuda,
|
|
config_field=config_field,
|
|
)
|
|
|
|
return env
|
|
|
|
|
|
def get_activation_fn(name: str):
|
|
"""Get activation function by name."""
|
|
if name == 'sin':
|
|
return SinActivation
|
|
elif name == 'tanh':
|
|
return torch.nn.Tanh
|
|
elif name == 'relu':
|
|
return torch.nn.ReLU
|
|
else:
|
|
raise ValueError(f"Unknown activation: {name}")
|
|
|
|
|
|
def evaluate_policy(model, env, n_episodes: int = 1):
|
|
"""
|
|
Evaluate the trained policy.
|
|
|
|
Returns:
|
|
Dictionary of evaluation metrics
|
|
"""
|
|
episode_rewards = []
|
|
episode_data = []
|
|
|
|
for episode in range(n_episodes):
|
|
obs, info = env.reset()
|
|
done = False
|
|
episode_reward = 0
|
|
steps = 0
|
|
|
|
ep_data = {
|
|
'actions': [],
|
|
'observations': [],
|
|
'rewards': [],
|
|
'cd': [],
|
|
'cl': [],
|
|
}
|
|
|
|
while not done:
|
|
action, _states = model.predict(obs, deterministic=True)
|
|
obs, reward, terminated, truncated, info = env.step(action)
|
|
done = terminated or truncated
|
|
|
|
episode_reward += reward
|
|
steps += 1
|
|
|
|
# Record data
|
|
ep_data['actions'].append(action)
|
|
ep_data['observations'].append(obs)
|
|
ep_data['rewards'].append(reward)
|
|
ep_data['cd'].append(info.get('cd', 0))
|
|
ep_data['cl'].append(info.get('cl', 0))
|
|
|
|
episode_rewards.append(episode_reward)
|
|
episode_data.append(ep_data)
|
|
|
|
print(f" Episode {episode + 1}/{n_episodes}: "
|
|
f"Reward = {episode_reward:.2f}, "
|
|
f"Steps = {steps}, "
|
|
f"Avg CD = {np.mean(ep_data['cd']):.4f}")
|
|
|
|
return {
|
|
'mean_reward': np.mean(episode_rewards),
|
|
'std_reward': np.std(episode_rewards),
|
|
'episodes': episode_data,
|
|
}
|
|
|
|
|
|
def main():
|
|
"""Main training loop."""
|
|
args = parse_args()
|
|
|
|
print("=" * 70)
|
|
print(f"DynamisLab - CFD Flow Control Training")
|
|
print(f"Run: {args.run_name}")
|
|
print("=" * 70)
|
|
|
|
# Create environment
|
|
print(f"\n[1/4] Creating CFD environment on GPU:{args.device_id}...")
|
|
env = create_env(args.device_id)
|
|
print(f" Action space: {env.action_space}")
|
|
print(f" Observation space: {env.observation_space}")
|
|
|
|
# Create or load model
|
|
print(f"\n[2/4] Setting up PPO model...")
|
|
device = torch.device(f"cuda:{args.cuda_device}")
|
|
|
|
if args.resume:
|
|
print(f" Resuming from: {args.resume}")
|
|
model = PPO.load(args.resume, env=env, device=device)
|
|
else:
|
|
activation_fn = get_activation_fn(args.activation)
|
|
model = PPO(
|
|
"MlpPolicy",
|
|
env=env,
|
|
learning_rate=args.learning_rate,
|
|
n_steps=args.n_steps,
|
|
batch_size=args.batch_size,
|
|
gamma=args.gamma,
|
|
policy_kwargs=dict(activation_fn=activation_fn),
|
|
device=device,
|
|
verbose=1,
|
|
)
|
|
|
|
print(f" Activation: {args.activation}")
|
|
print(f" Device: {device}")
|
|
print(f" Learning rate: {args.learning_rate}")
|
|
print(f" Steps per iteration: {args.n_steps}")
|
|
print(f" Batch size: {args.batch_size}")
|
|
|
|
# Setup logging
|
|
tensorboard_dir = get_tensorboard_logdir(args.run_name)
|
|
writer = SummaryWriter(log_dir=str(tensorboard_dir))
|
|
print(f" TensorBoard: {tensorboard_dir}")
|
|
|
|
# Training loop
|
|
print(f"\n[3/4] Training for {args.total_timesteps} iterations...")
|
|
|
|
best_reward = -np.inf
|
|
history_data = []
|
|
|
|
for iteration in range(args.total_timesteps):
|
|
# Train
|
|
model.learn(total_timesteps=args.n_steps, reset_num_timesteps=False)
|
|
|
|
# Evaluate
|
|
print(f"\n--- Iteration {iteration + 1}/{args.total_timesteps} ---")
|
|
eval_results = evaluate_policy(model, env, n_episodes=args.eval_episodes)
|
|
|
|
mean_reward = eval_results['mean_reward']
|
|
std_reward = eval_results['std_reward']
|
|
|
|
# Log to TensorBoard
|
|
writer.add_scalar('eval/mean_reward', mean_reward, iteration)
|
|
writer.add_scalar('eval/std_reward', std_reward, iteration)
|
|
|
|
# Extract CD/CL from last episode
|
|
if len(eval_results['episodes']) > 0:
|
|
last_ep = eval_results['episodes'][-1]
|
|
avg_cd = np.mean(last_ep['cd'])
|
|
avg_cl = np.mean(last_ep['cl'])
|
|
writer.add_scalar('eval/avg_cd', avg_cd, iteration)
|
|
writer.add_scalar('eval/avg_cl', avg_cl, iteration)
|
|
|
|
# Save best model
|
|
if mean_reward > best_reward:
|
|
best_reward = mean_reward
|
|
model_path = get_model_path(f"{args.run_name}_best")
|
|
model.save(str(model_path))
|
|
print(f" ✓ New best model saved: {model_path} (reward: {mean_reward:.2f})")
|
|
|
|
# Periodic save
|
|
if (iteration + 1) % args.save_freq == 0:
|
|
model_path = get_model_path(f"{args.run_name}_iter{iteration + 1}")
|
|
model.save(str(model_path))
|
|
print(f" Checkpoint saved: {model_path}")
|
|
|
|
# Store history
|
|
history_data.append(eval_results['episodes'])
|
|
|
|
# Final evaluation
|
|
print(f"\n[4/4] Final evaluation...")
|
|
final_results = evaluate_policy(model, env, n_episodes=5)
|
|
print(f" Final mean reward: {final_results['mean_reward']:.2f} ± {final_results['std_reward']:.2f}")
|
|
|
|
# Save final model and history
|
|
final_model_path = get_model_path(f"{args.run_name}_final")
|
|
model.save(str(final_model_path))
|
|
print(f" Final model saved: {final_model_path}")
|
|
|
|
history_path = get_output_path(f"{args.run_name}_history.pkl")
|
|
with open(history_path, 'wb') as f:
|
|
pickle.dump(history_data, f)
|
|
print(f" Training history saved: {history_path}")
|
|
|
|
# Cleanup
|
|
writer.close()
|
|
env.close()
|
|
|
|
print("\n" + "=" * 70)
|
|
print("Training complete!")
|
|
print("=" * 70)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|