From b3ee72e1445a1e56d6689e422af311111d6eedcc Mon Sep 17 00:00:00 2001 From: Frank14f <1515444314@qq.com> Date: Wed, 1 Jul 2026 20:10:27 +0800 Subject: [PATCH] =?UTF-8?q?feat(train):=20V5=20parameterized=20training=20?= =?UTF-8?q?pipeline=20=E2=80=94=20Karman=20+=20Illusion=20verified?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Calibration-driven, no_bias only, 2000x600 grid. All cases share unified env/train/calibrate pattern. Multi-GPU server deployment ready. Core additions: - calibrate.py: Phase 0 calibration (karman/illusion), produces calibration.json with rounded FORCE_SCALE, SENS_SCALE, SIM_BP/VAL - env_karman.py: parameterized Karman cloak env (calibration + config_path) - env_illusion.py: illusion env with FFT harmonics target (S_DIM=14) - env_vortex.py: vortex cloaking env (lamb/taylor, MAX_STEPS=150) - train_karman.py, train_illusion.py: parameterized training scripts - launch_multi.sh: sequential multi-GPU launcher (7-min staggered) - SERVER_DEPLOY.md: complete server setup, calibration, training guide - calibrations/re100/ & calibrations/illusion_1L/: pre-run calibrations Fixes: - SIM_VAL[-1] 0.95 -> 1.0 (r_sim maps to full [0,1] range) - Cross-Re configs: re50/200/400 (viscosity-only variants) Verified end-to-end on GPU0+GPU1: - Karman V5 20-ep: best reward 0.459 at Ep16 (monotonic rise) - Illusion 20-ep: best reward 0.224 at Ep19 (harmonics, DTW learning) Co-authored-by: Cursor --- configs/config_lbm_karman_2000x600_re200.json | 50 ++ configs/config_lbm_karman_2000x600_re400.json | 50 ++ configs/config_lbm_karman_2000x600_re50.json | 50 ++ src/drl_pinball/train/SERVER_DEPLOY.md | 189 ++++++ src/drl_pinball/train/TRAIN_KNOWLEDGE.md | 177 +++++- src/drl_pinball/train/calibrate.py | 573 ++++++++++++++++++ .../calibrations/illusion_1L/calibration.json | 50 ++ .../illusion_1L/target_harmonics.json | 194 ++++++ .../train/calibrations/re100/calibration.json | 49 ++ src/drl_pinball/train/env_illusion.py | 426 +++++++++++++ src/drl_pinball/train/env_karman.py | 419 +++++++++++++ src/drl_pinball/train/env_karman_2000x600.py | 12 +- src/drl_pinball/train/env_vortex.py | 378 ++++++++++++ src/drl_pinball/train/launch_multi.sh | 147 +++++ src/drl_pinball/train/train_illusion.py | 207 +++++++ src/drl_pinball/train/train_karman.py | 226 +++++++ .../train/visualize_and_analyze.py | 492 +++++++++++++++ 17 files changed, 3652 insertions(+), 37 deletions(-) create mode 100644 configs/config_lbm_karman_2000x600_re200.json create mode 100644 configs/config_lbm_karman_2000x600_re400.json create mode 100644 configs/config_lbm_karman_2000x600_re50.json create mode 100644 src/drl_pinball/train/SERVER_DEPLOY.md create mode 100644 src/drl_pinball/train/calibrate.py create mode 100644 src/drl_pinball/train/calibrations/illusion_1L/calibration.json create mode 100644 src/drl_pinball/train/calibrations/illusion_1L/target_harmonics.json create mode 100644 src/drl_pinball/train/calibrations/re100/calibration.json create mode 100644 src/drl_pinball/train/env_illusion.py create mode 100644 src/drl_pinball/train/env_karman.py create mode 100644 src/drl_pinball/train/env_vortex.py create mode 100755 src/drl_pinball/train/launch_multi.sh create mode 100644 src/drl_pinball/train/train_illusion.py create mode 100644 src/drl_pinball/train/train_karman.py create mode 100644 src/drl_pinball/train/visualize_and_analyze.py diff --git a/configs/config_lbm_karman_2000x600_re200.json b/configs/config_lbm_karman_2000x600_re200.json new file mode 100644 index 0000000..3f31416 --- /dev/null +++ b/configs/config_lbm_karman_2000x600_re200.json @@ -0,0 +1,50 @@ +{ + "_doc": "Karman Cloak Re200: uniform inlet, free-slip walls, 2000x600 grid. Pinball centered.", + "grid": { + "lattice_model": "D2Q9", + "nx": 2000, + "ny": 600, + "nz": 1 + }, + "physics": { + "data_type": "FP32", + "viscosity": 0.002, + "velocity": 0.01, + "rho": 1.0 + }, + "method": { + "collision": "MRT", + "streaming": "double_buffer", + "store_precision": "FP32", + "ddf_shifting": false, + "les": { + "enabled": false, + "cs": 0.16, + "closed_form": true + }, + "trt": { + "magic_param": 0.1875 + }, + "inlet": { + "profile": "uniform", + "scheme": "regularized", + "trt_neq_damp": 0.5, + "regularized_neq_damp": 0.5 + }, + "outlet": { + "mode": "neq_extrap", + "backflow_clamp": true, + "blend_alpha": 0.7, + "srt_neq_damp": 0.5 + }, + "y_wall_bc": "free_slip", + "omega_guard": { + "min": 0.01, + "max": 1.99 + } + }, + "cuda": { + "threads_per_block": 256, + "compute_capability": "auto" + } +} diff --git a/configs/config_lbm_karman_2000x600_re400.json b/configs/config_lbm_karman_2000x600_re400.json new file mode 100644 index 0000000..a5e1df1 --- /dev/null +++ b/configs/config_lbm_karman_2000x600_re400.json @@ -0,0 +1,50 @@ +{ + "_doc": "Karman Cloak Re400: uniform inlet, free-slip walls, 2000x600 grid. Pinball centered.", + "grid": { + "lattice_model": "D2Q9", + "nx": 2000, + "ny": 600, + "nz": 1 + }, + "physics": { + "data_type": "FP32", + "viscosity": 0.001, + "velocity": 0.01, + "rho": 1.0 + }, + "method": { + "collision": "MRT", + "streaming": "double_buffer", + "store_precision": "FP32", + "ddf_shifting": false, + "les": { + "enabled": false, + "cs": 0.16, + "closed_form": true + }, + "trt": { + "magic_param": 0.1875 + }, + "inlet": { + "profile": "uniform", + "scheme": "regularized", + "trt_neq_damp": 0.5, + "regularized_neq_damp": 0.5 + }, + "outlet": { + "mode": "neq_extrap", + "backflow_clamp": true, + "blend_alpha": 0.7, + "srt_neq_damp": 0.5 + }, + "y_wall_bc": "free_slip", + "omega_guard": { + "min": 0.01, + "max": 1.99 + } + }, + "cuda": { + "threads_per_block": 256, + "compute_capability": "auto" + } +} diff --git a/configs/config_lbm_karman_2000x600_re50.json b/configs/config_lbm_karman_2000x600_re50.json new file mode 100644 index 0000000..9eb9176 --- /dev/null +++ b/configs/config_lbm_karman_2000x600_re50.json @@ -0,0 +1,50 @@ +{ + "_doc": "Karman Cloak Re50: uniform inlet, free-slip walls, 2000x600 grid. Pinball centered.", + "grid": { + "lattice_model": "D2Q9", + "nx": 2000, + "ny": 600, + "nz": 1 + }, + "physics": { + "data_type": "FP32", + "viscosity": 0.008, + "velocity": 0.01, + "rho": 1.0 + }, + "method": { + "collision": "MRT", + "streaming": "double_buffer", + "store_precision": "FP32", + "ddf_shifting": false, + "les": { + "enabled": false, + "cs": 0.16, + "closed_form": true + }, + "trt": { + "magic_param": 0.1875 + }, + "inlet": { + "profile": "uniform", + "scheme": "regularized", + "trt_neq_damp": 0.5, + "regularized_neq_damp": 0.5 + }, + "outlet": { + "mode": "neq_extrap", + "backflow_clamp": true, + "blend_alpha": 0.7, + "srt_neq_damp": 0.5 + }, + "y_wall_bc": "free_slip", + "omega_guard": { + "min": 0.01, + "max": 1.99 + } + }, + "cuda": { + "threads_per_block": 256, + "compute_capability": "auto" + } +} diff --git a/src/drl_pinball/train/SERVER_DEPLOY.md b/src/drl_pinball/train/SERVER_DEPLOY.md new file mode 100644 index 0000000..fc0b110 --- /dev/null +++ b/src/drl_pinball/train/SERVER_DEPLOY.md @@ -0,0 +1,189 @@ +# Server Deployment Guide + +## Prerequisites + +- Linux with NVIDIA GPU(s) and CUDA toolkit installed +- Python 3.10+ +- CelerisLab submodule + +## One-Time Setup + +```bash +# 1. Clone with submodules +git clone --recurse-submodules +cd DynamisLab + +# 2. Create conda environment +conda create -n pycuda_3_10 python=3.10 +conda activate pycuda_3_10 + +# 3. Install PyCUDA (match your CUDA version) +pip install pycuda + +# 4. Install CelerisLab +pip install -e CelerisLab + +# 5. Install DynamisLab +pip install -e . + +# 6. Verify +python -c "from CelerisLab import Simulation; print('OK')" +``` + +## Calibration (one per case, before training) + +Every case MUST be calibrated once. This produces calibration.json + target files. + +### Karman Cloak (Re100, SI=800) +```bash +cd src/drl_pinball/train +conda run -n pycuda_3_10 python calibrate.py \ + --case re100 --device-id 0 \ + --config ../../../configs/config_lbm_karman_2000x600.json +``` + +### Karman Cloak (Re200, SI=500) +```bash +conda run -n pycuda_3_10 python calibrate.py \ + --case re200 --device-id 0 --si 500 \ + --config ../../../configs/config_lbm_karman_2000x600_re200.json +``` + +### Karman Cloak (Re50, SI=1600 / Re400, SI=400) +```bash +conda run -n pycuda_3_10 python calibrate.py \ + --case re50 --device-id 0 --si 1600 \ + --config ../../../configs/config_lbm_karman_2000x600_re50.json + +conda run -n pycuda_3_10 python calibrate.py \ + --case re400 --device-id 0 --si 400 \ + --config ../../../configs/config_lbm_karman_2000x600_re400.json +``` + +### Illusion (1L target, SI=600) +```bash +conda run -n pycuda_3_10 python calibrate.py \ + --case illusion_1L --device-id 0 --si 600 --scene illusion \ + --config ../../../configs/config_lbm_karman_2000x600.json +``` + +## Training + +### CRITICAL: Sequential GPU Startup + +Each GPU needs ~7 minutes between starts because CelerisLab compiles CUDA kernels +during the first Simulation() constructor. If two start simultaneously, the kernel +cache gets corrupted → reward=0.000 forever. + +### Single-GPU Training (for testing) + +```bash +# Karman Re100, 20 episodes +conda run -n pycuda_3_10 python -u train_karman.py \ + --case-name re100_test --device-id 0 --seed 42 \ + --config ../../../configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/re100/calibration.json \ + --total-episodes 20 + +# Illusion 1L, 20 episodes +conda run -n pycuda_3_10 python -u train_illusion.py \ + --case-name illusion_test --device-id 0 --seed 42 \ + --config ../../../configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/illusion_1L/calibration.json \ + --total-episodes 20 +``` + +### Multi-GPU Training (6 GPUs, using launch script) + +```bash +# Re100 Karman, 6 seeds +bash launch_multi.sh \ + --case-name re100_karman --seeds 42,43,44,45,46,47 \ + --gpus 0,1,2,3,4,5 --episodes 500 \ + --config ../../../configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/re100/calibration.json + +# Re200 Karman, transfer learning from Re100 +bash launch_multi.sh \ + --case-name re200_karman --seeds 42,43,44,45,46,47 \ + --gpus 0,1,2,3,4,5 --episodes 500 \ + --config ../../../configs/config_lbm_karman_2000x600_re200.json \ + --calibration calibrations/re200/calibration.json \ + --transfer output/re100_karman_seed42/models/best_model.zip +``` + +### Manual Multi-GPU (tmux/screen) + +If launch_multi.sh doesn't work, start manually with delays: + +```bash +# GPU 0 +nohup conda run -n pycuda_3_10 python -u train_karman.py \ + --case-name re100_karman --device-id 0 --seed 42 \ + --config ../../../configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/re100/calibration.json \ + --total-episodes 500 > output/re100_karman_seed42/nohup.log 2>&1 & + +# WAIT 7 minutes (until "Env ready" appears in log) +sleep 420 + +# GPU 1 +nohup conda run -n pycuda_3_10 python -u train_karman.py \ + --case-name re100_karman --device-id 1 --seed 43 \ + --config ../../../configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/re100/calibration.json \ + --total-episodes 500 > output/re100_karman_seed43/nohup.log 2>&1 & + +# Repeat for GPU 2,3,4,5 with seeds 44,45,46,47 +``` + +## Monitoring + +```bash +# Check training progress +tail -f output/re100_karman_seed42/train.log + +# TensorBoard +tensorboard --logdir output/re100_karman_seed42/tb --port 6006 --bind_all + +# Monitor GPU usage +watch -n 1 nvidia-smi +``` + +## Output Structure + +``` +output/ +└── re100_karman_seed42/ + ├── models/ + │ ├── best_model.zip # Best reward model + │ ├── final_model.zip # Final iteration model + │ └── chkpt_ep*.zip # Checkpoints every 10 episodes + ├── tb/ # TensorBoard logs + ├── train.log # Training log + ├── calibration.json # Copy of calibration used + ├── vec_normalize.pkl # VecNormalize statistics + └── meta.json # Run metadata +``` + +## Stopping + +```bash +# Stop all training +pkill -f train_karman +pkill -f train_illusion + +# Or by PID +ps aux | grep train_karman | grep -v grep | awk '{print $2}' | xargs kill +``` + +## Troubleshooting + +| Symptom | Cause | Fix | +|---------|-------|-----| +| reward=0.000 forever | Kernel compilation race (two inits at once) | Kill both, `rm -f ~/CelerisLab/src/CelerisLab/lbm/kernels/config/config_objects.h ~/CelerisLab/src/CelerisLab/lbm/kernels/kernel.ptx`, restart sequentially | +| reward=NaN | FORCE_SCALE too small | Check calibration.json, re-run calibrate.py | +| reward flat at Stage0 | Action not being applied | Check `_action_to_omega` sign, `set_body` call | +| CUDA OOM | PyTorch+PyCUDA memory conflict | Reduce `--batch-size 32` | +| ImportError: CelerisLab | Not installed | Run `pip install -e CelerisLab` in conda env | +| Conda env not found | Wrong environment | Use `conda env list` to verify `pycuda_3_10` exists | diff --git a/src/drl_pinball/train/TRAIN_KNOWLEDGE.md b/src/drl_pinball/train/TRAIN_KNOWLEDGE.md index 3a7a4f0..9216d9f 100644 --- a/src/drl_pinball/train/TRAIN_KNOWLEDGE.md +++ b/src/drl_pinball/train/TRAIN_KNOWLEDGE.md @@ -1,8 +1,7 @@ -# Karman Cloak Training — Knowledge Document +# Karman Cloak Training — Knowledge Document (V5) -> **For new developers taking over**: This document contains everything you need to know -> to run, debug, and improve the DRL training pipeline for Karman Cloak on the new -> CelerisLab solver. Read this before touching any code. +> **V5 (2026-07-01)**: Parameterized, calibration-driven, no_bias only, multi-GPU ready. +> Original V4 files preserved as `env_karman_2000x600.py` and `train_karman_2000x600.py`. --- @@ -13,36 +12,145 @@ hydrodynamic cloaking — making the downstream flow match the "undisturbed" flo (as if the pinball weren't there). The upstream disturbance cylinder generates a Kármán vortex street; the pinball must cancel it. -- **CFD**: CelerisLab LBM solver, D2Q9 MRT, 2000×600 grid, uniform inlet, free-slip walls -- **DRL**: PPO with Sin activation, 64×64 MLP, SB3 + VecNormalize -- **Two training modes**: Bias (with [0,-4,4] action offset) and NoBias (scale=12, no offset) +- **CFD**: CelerisLab LBM solver, D2Q9 MRT, 2000x600 grid, uniform inlet, free-slip walls +- **DRL**: PPO with Sin activation, 64x64 MLP, SB3 + VecNormalize +- **V5 mode**: No_bias only (ACTION_SCALE=12, ACTION_BIAS=[0,0,0]). All cases use calibration-first workflow. + +## 0. V5 Quick Start (Server Deployment) + +```bash +# 1. Calibrate (once per case, ~5 min on single GPU) +cd src/drl_pinball/train +conda run -n pycuda_3_10 python calibrate.py \ + --case re100 --device-id 0 \ + --config configs/config_lbm_karman_2000x600.json + +# 2. Launch multi-seed training on server (sequential, 7 min between GPUs) +bash launch_multi.sh \ + --case-name re100_karman --seeds 42,43,44,45,46,47 \ + --gpus 0,1,2,3,4,5 --episodes 500 \ + --config configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/re100/calibration.json + +# 3. Transfer learning to another Re +conda run -n pycuda_3_10 python calibrate.py \ + --case re200 --device-id 0 --si 500 \ + --config configs/config_lbm_karman_2000x600_re200.json + +bash launch_multi.sh \ + --case-name re200_karman --seeds 42,43,44 --gpus 0,1,2 \ + --config configs/config_lbm_karman_2000x600_re200.json \ + --calibration calibrations/re200/calibration.json \ + --transfer output/re100_karman_seed42/models/best_model.zip + +# 4. Monitor +tail -f output/re100_karman_seed42/train.log +tensorboard --logdir output/re100_karman_seed42/tb --port 6006 +``` --- -## 2. File Structure (Clean) +## 2. File Structure (V5) ``` train/ ├── __init__.py -├── env_karman_2000x600.py # Gym environment (final V4) -├── train_karman_2000x600.py # PPO training script (final V4) +│ +├── # V5 ACTIVE FILES (parameterized, no_bias, calibration-driven) +│ +├── calibrate.py # Phase 0 calibration (produces calibration.json + target.npy) +├── env_karman.py # Parameterized Karman cloak env (loads calibration JSON) +├── env_vortex.py # Vortex cloak env (lamb/taylor, MAX_STEPS=150) +├── train_karman.py # Parameterized training script (--calibration, --config) +├── launch_multi.sh # Sequential multi-GPU server launcher +│ +├── # V5 SUPPORT FILES (unchanged from V4) +│ ├── symmetry_wrapper.py # G-mirror symmetry augmentation (per-rollout) -├── phase0_baseline_measure.py # Stage0/Bias baseline measurement tool -├── analyze_final.py # Training curve plotting & degradation analysis -├── TRAIN_KNOWLEDGE.md # This file -├── target.npy # Pre-recorded target signal (reusable) -├── nohup_bias.log # nohup stdout for current Bias run -├── nohup_nobias.log # nohup stdout for current NoBias run -└── output/ - ├── bias_seed42_s2048_e10_v4/ # Current V4 Bias training (running) - ├── nobias_seed42_s2048_e10_v4/ # Current V4 NoBias training (running) - ├── stage_baseline_2000x600/ # Phase 0 baseline data (reference) - └── degradation_diag/ # Analysis plots from debugging +├── phase0_baseline_measure.py # Legacy baseline measurement tool (reference) +├── visualize_and_analyze.py # Flow-field visualization & analysis +│ +├── # V4 BACKUP FILES (preserved, NOT active) +│ +├── env_karman_2000x600.py # Original V4 env (hardcoded FORCE_SCALE, bias support) +├── train_karman_2000x600.py # Original V4 training (--no-bias flag) +├── analyze_final.py # V3 training curve plotting (legacy paths) +│ +├── # Calibration & output +│ +├── calibrations/ # Per-case calibration files +│ └── re100/ +│ ├── calibration.json # FORCE_SCALE, SENS_SCALE, SIM_BP/VAL, etc. +│ ├── target.npy # Target sensor signals (FIFO_LEN, 6) +│ └── calibrate.log # Calibration run log +│ +└── output/ # Training outputs + └── re100_karman_seed42/ + ├── models/ # best_model.zip, final_model.zip, chkpt_ep*.zip + ├── tb/ # TensorBoard logs + ├── train.log # Training run log + ├── calibration.json # Copy of calibration used + ├── vec_normalize.pkl # VecNormalize statistics + └── meta.json # Run metadata ``` --- -## 3. How to Start Training +## 3. Calibration Workflow (V5 — ALWAYS RUN FIRST) + +Every case MUST run `calibrate.py` before training. This produces: + +- `calibration.json`: FORCE_SCALE, SENS_SCALE, dtw_norm_scale, SIM_BP, SIM_VAL, reward constants +- `target.npy`: Target sensor signals (150 steps x 6 channels, legacy-equiv) + +The calibration measures Stage0 (zero rotation) and Stage1 (open-loop reference action +equivalent to legacy [0,-4,4]*U0 bias) to compute: +- FORCE_SCALE = combined max|force| across both stages +- SENS_SCALE = combined max|sensor| (legacy-equiv) across both stages +- dtw_norm_scale = target's 3 uy-channel avg std +- SIM_BP/SIM_VAL = piecewise DTW mapping from Stage0/Stage1 sim mean values + +Calibration is IMMUTABLE — once produced, never modify. Training and inference use it as-is. + +### Cross-Re SI guidance + +Based on ~18 samples per vortex shedding cycle: +| Case | SI | Rationale | +|------|----|-----------| +| re50 | 1600 | Lower Re, longer period | +| re100 | 800 | Verified (~19 samples/cycle) | +| re200 | 500 | ~18 samples/cycle | +| re400 | 400 | ~18 samples/cycle | + +--- + +## 4. Calibration Results (re100) + +``` +FORCE_SCALE = 0.002429 +SENS_SCALE = 0.7543 +dtw_norm_scale = 0.2043 +Stage0 sim_mean = 0.3166 (-> SIM_VAL=0.2) +Stage1 sim_mean = 0.8158 (-> SIM_VAL=0.5) +SIM_BP = [0.0, 0.3166, 0.8158, 0.8895, 0.9448, 1.0] +SIM_VAL = [0.0, 0.2, 0.5, 0.8, 0.9, 1.0] +``` + +Verified: env with calibration gives zero-action reward = 0.076 (Stage0 level, correct). + +--- + +## 5. Cross-Re Configs + +All share the same 2000x600 grid, uniform inlet, free-slip walls. Only viscosity differs: +- `configs/config_lbm_karman_2000x600.json` (v=0.004, Re=100) +- `configs/config_lbm_karman_2000x600_re50.json` (v=0.008, Re=50) +- `configs/config_lbm_karman_2000x600_re200.json` (v=0.002, Re=200) +- `configs/config_lbm_karman_2000x600_re400.json` (v=0.001, Re=400) + +--- + +## 6. Original "How to Start Training" (V4 — kept for reference) ### Prerequisites - conda env `pycuda_3_10` with PyCUDA, PyTorch, SB3, tensorboard @@ -409,12 +517,19 @@ balances) — only use if you re-record the target with the same inlet scheme. ## 13. Future Work -1. **Multi-seed training**: Run 3-5 seeds to show variance bands in paper -2. **Illusion adaptation**: Same pipeline, different target (target cylinder wake). - Need to re-record target, recompute DTW norm_scale, adjust sim breakpoints. -3. **Steady cloak**: Simpler case (no upstream disturbance cylinder). Target is - uniform flow. DTW sim may need different handling (target std≈0). -4. **Symmetry ablation study**: Compare with/without G-mirror to quantify benefit -5. **Longer training**: 500 episodes may not be enough for NoBias to fully learn - r_cl. Try 1000 episodes. -6. **Learning rate schedule**: Try lr decay after peak to prevent degradation +1. **Re50/re200/re400 transfer learning**: Calibrate each Re, then transfer from re100 best model using `--transfer` flag. Use adjusted SI per Re. +2. **Vortex cloak**: `env_vortex.py` ready. Transfer from re100 model. Lamb and Taylor variants. +3. **Illusion**: `env_illusion.py` needed. S_DIM=14, harmonics target reconstruction. Adjusted pinball/sensor positions. +4. **Steady cloak**: Simpler case (no upstream disturbance cylinder). Target is uniform flow. +5. **Symmetry ablation study**: Compare with/without G-mirror to quantify benefit. +6. **Longer training**: 500 episodes may not be enough. Try 1000 episodes. +7. **Learning rate schedule**: Try lr decay after peak to prevent degradation. + +## 14. V5 Design Principles + +1. **Calibration-first**: Every case runs calibrate.py before training. Produces a single immutable JSON. +2. **No-bias only**: action_scale=12, action_bias=[0,0,0]. Simplifies G-mirror, removes human priors. +3. **Parameterized envs**: env_karman.py accepts calibration dict + config_path + si. No module-level hardcode. +4. **Sequential GPU startup**: launch_multi.sh enforces 7-min delay between launches. +5. **Immutable calibration**: Once produced, calibration is read-only for training and inference. Saved alongside model. +6. **r_sim maps to [0,1]**: SIM_VAL[-1]=1.0. Full-range normalized DTW similarity. diff --git a/src/drl_pinball/train/calibrate.py b/src/drl_pinball/train/calibrate.py new file mode 100644 index 0000000..dd67376 --- /dev/null +++ b/src/drl_pinball/train/calibrate.py @@ -0,0 +1,573 @@ +#!/usr/bin/env python3 +"""Phase 0 calibration: measure baselines and produce calibration.json for any case. + +Runs on new CelerisLab solver with 2000x600 config. Produces: + - target.npy (FIFO_LEN, 6) sensor signals (legacy-equiv) + - calibration.json FORCE_SCALE, SENS_SCALE, dtw_norm_scale, SIM_BP, SIM_VAL + +Workflow: + 1. Temp sim (dist_cyl + sensors only) -> warmup -> record target -> close + 2. Training sim (all 7 objects) -> warmup + 3. Stage0: zero-action measurement (150+100 steps) + 4. Stage1: open-loop reference measurement (legacy-equiv bias) + 5. Compute calibration constants from both stages + +Usage: + conda run -n pycuda_3_10 python calibrate.py \ + --case re100 --device-id 0 \ + --config configs/config_lbm_karman_2000x600.json +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from collections import deque +from pathlib import Path + +import numpy as np +import pycuda.driver as cuda; cuda.init() + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from CelerisLab import Simulation + +_CELERIS = Path("/home/frank14f/CelerisLab") +_CONFIG_H = _CELERIS / "src/CelerisLab/lbm/kernels/config/config_objects.h" +_PTX = _CELERIS / "src/CelerisLab/lbm/kernels/kernel.ptx" + +def _clean_cache(): + for p in [_CONFIG_H, _PTX]: + if p.exists(): + p.unlink() + +# --------------------------------------------------------------------------- +# Physics / geometry constants +# --------------------------------------------------------------------------- +L0 = 20.0 +U0 = 0.01 +RADIUS = L0 / 2.0 +NX = 2000 +NY = 600 +CENTER_Y = float(NY - 1) / 2.0 + +DIST_X = 600.0 +PINBALL_FRONT_X = 1000.0 +PINBALL_REAR_X = 1026.0 +SENSOR_X = 1200.0 + +FIFO_LEN = 150 +CONV_LEN = 30 +SENSOR_CC = 78.0 + +N_MEASURE = 100 + +# Reference action: legacy-equiv bias [0,-4,4]*U0, expressed as no_bias omega +# no_bias: omega = -(action * 12 + [0,0,0]) * U0 / R +# Target omega = [0, 0.004, -0.004] (= action_norm * 12 * U0 / R with action_norm [0, -4/12, 4/12]) +_REF_OMEGA = np.array([0.0, 0.004, -0.004], dtype=np.float32) + +# Reward constants +K_CD = 50.0; K_CL = 100.0 +W_CD = 0.30; W_CL = 0.30; W_SIM = 0.40 +FLOOR_CD = 0.10; FLOOR_CL = 0.10; FLOOR_SIM = 0.10 +FLOOR_PENALTY = 0.05 + +ACTION_SCALE = 12.0 +ACTION_BIAS = np.array([0.0, 0.0, 0.0], dtype=np.float32) + + +# --------------------------------------------------------------------------- +# DTW utilities +# --------------------------------------------------------------------------- +def calc_lag(target: np.ndarray, state: np.ndarray) -> int: + t_mean = np.mean(target); s_mean = np.mean(state) + corr = np.correlate(target - t_mean, state - s_mean, mode="full") + lags = np.arange(-len(target) + 1, len(target)) + return int(lags[np.argmax(corr)]) + + +def calc_dtw_sim(target: np.ndarray, state: np.ndarray, + norm_scale: float = 1.0) -> float: + n, m = len(target), len(state) + dtw = np.full((n + 1, m + 1), np.inf); dtw[0, 0] = 0.0 + for i in range(1, n + 1): + for j in range(1, m + 1): + cost = abs(float(target[i - 1]) - float(state[j - 1])) + last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1]) + dtw[i, j] = cost + last_min + raw = 1.0 - dtw[n, m] / (float(n) * norm_scale) + return float(max(0.0, raw)) + + +def compute_similarity(target_states, fifo_states, conv_len, norm_scale): + target = np.asarray(target_states, dtype=np.float64) + state = np.asarray(fifo_states, dtype=np.float64) + id_sens = 3 + target_seq = target[conv_len:2 * conv_len, id_sens] + state_seq = state[-conv_len:, id_sens] + lag = calc_lag(target_seq, state_seq) + sim = 0.0 + for i in range(6): + t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len] + s_seq = state[-conv_len:, i] + sim += calc_dtw_sim(t_seq, s_seq, norm_scale=norm_scale) + return float(sim / 6.0) + + +# --------------------------------------------------------------------------- +# CtxGuard +# --------------------------------------------------------------------------- +class CtxGuard: + def __init__(self, sim): + self.sim = sim + def __enter__(self): + if self.sim is not None: + self.sim.ctx._ctx.push() + def __exit__(self, *exc): + if self.sim is not None: + self.sim.ctx._ctx.pop() + return False + + +def gpu_block(sim, fn): + with CtxGuard(sim): + fn() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def read_obs(sim, dist_id, sensor_ids, pinball_ids) -> np.ndarray: + obs = [] + if dist_id is not None: + obs = list(sim.read_force(dist_id, normalize=True)) + else: + obs = [0.0, 0.0] # no disturbance cylinder (illusion) + for sid in sensor_ids: + s = sim.read_sensor(sid, normalize=True) + obs.extend([float(s[0]), float(s[1])]) + for pid in pinball_ids: + obs.extend(sim.read_force(pid, normalize=True)) + return np.array(obs, dtype=np.float32) + + +def set_omega(sim, pinball_ids, omega): + for pid, w in zip(pinball_ids, omega): + sim.set_body(pid, omega=float(w)) + + +def run_stage(sim, dist_id, sensor_ids, pinball_ids, target_states, + omega_vec, dtw_norm_scale, si, n_measure=N_MEASURE): + gpu_block(sim, lambda: set_omega(sim, pinball_ids, omega_vec)) + fifo = deque(maxlen=FIFO_LEN) + for _ in range(FIFO_LEN): + gpu_block(sim, lambda: sim.run(si, zero_obs=True)) + obs = read_obs(sim, dist_id, sensor_ids, pinball_ids) + sl = obs[2:14].copy() + sl[0:6] *= SENSOR_CC + fifo.append(sl) + + obs_slices, sims = [], [] + for _ in range(n_measure): + gpu_block(sim, lambda: sim.run(si, zero_obs=True)) + obs = read_obs(sim, dist_id, sensor_ids, pinball_ids) + sl = obs[2:14].copy() + sl[0:6] *= SENSOR_CC + fifo.append(sl) + obs_slices.append(sl.copy()) + sim_val = compute_similarity(target_states, np.array(list(fifo)), + CONV_LEN, dtw_norm_scale) + sims.append(sim_val) + + return { + "obs_slices": np.array(obs_slices, dtype=np.float32), + "sims": np.array(sims, dtype=np.float32), + } + + +# --------------------------------------------------------------------------- +# Illusion calibration +# --------------------------------------------------------------------------- + +# Illusion-specific geometry +_ILL_PINBALL_FRONT_X = 380.0 +_ILL_PINBALL_REAR_X = 406.0 +_ILL_SENSOR_X = 600.0 +_ILL_TARGET_X = 400.0 +_ILL_TARGET_RADIUS = 1.0 * L0 +_ILL_REF_OMEGA = np.array([0.0, 0.002, -0.002], dtype=np.float32) # legacy-equiv bias + +def _analyze_harmonics_for_calib(states, n_harmonics=5): + N, D = states.shape + result = [] + for d in range(D): + y = states[:, d] + fft_coef = np.fft.rfft(y) + freqs = np.fft.rfftfreq(N, d=1) + amps = 2 * np.abs(fft_coef) / N + phases = np.angle(fft_coef) + idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1 + harmonics = { + 'dc': float(np.real(fft_coef[0]) / N), + 'amps': amps[idx].tolist(), + 'freqs': freqs[idx].tolist(), + 'phases': phases[idx].tolist(), + } + result.append(harmonics) + return result + + +def _calibrate_illusion(case, config_path, device_id, si, out_dir, log, warmup): + # ---- Step 1: Record target (target cylinder + 3 sensors) ---- + log("Step 1: Recording illusion target (target cyl + sensors)...") + _clean_cache() + sim = Simulation(lbm_config_path=config_path, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + sim.add_body("circle", center=(_ILL_TARGET_X, CENTER_Y, 0.0), + radius=_ILL_TARGET_RADIUS) + s0 = sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0) + s1 = sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y, 0.0), radius=5.0) + s2 = sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0) + sim.initialize() + gpu_block(sim, lambda: sim.run(warmup, zero_obs=True)) + log(f" Warmup done ({warmup} steps)") + + target_states = np.zeros((FIFO_LEN, 8), dtype=np.float32) + for i in range(FIFO_LEN): + gpu_block(sim, lambda: sim.run(si, zero_obs=True)) + target_states[i] = [ + sim.read_sensor(s0, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s0, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s1, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s1, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s2, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s2, normalize=True)[1] * SENSOR_CC, + sim.read_force(0, normalize=True)[0], + sim.read_force(0, normalize=True)[1], + ] + sim.close() + # Sensor-only portion for DTW + target_sensor = target_states[:, 0:6].copy() + np.save(str(out_dir / "target.npy"), target_sensor) + + # Harmonics + harmonics = _analyze_harmonics_for_calib(target_states, n_harmonics=5) + with open(out_dir / "target_harmonics.json", "w") as f: + json.dump(harmonics, f, indent=2) + + target_std = np.std(target_sensor, axis=0) + uy_std_avg = float(np.mean([target_std[1], target_std[3], target_std[5]])) + dtw_norm_scale = max(uy_std_avg, 0.01) + log(f" Target recorded. DTW norm_scale: {dtw_norm_scale:.4f}") + + # ---- Step 2: Training sim (3 sensors + 3 pinball) ---- + log("Step 2: Creating training sim (6 objects)...") + _clean_cache() + sim = Simulation(lbm_config_path=config_path, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + sensor_ids = [ + sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0), + sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y, 0.0), radius=5.0), + sim.add_body("sensor", center=(_ILL_SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0), + ] + sim.add_body("circle", center=(_ILL_PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS) + sim.add_body("circle", center=(_ILL_PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS) + sim.add_body("circle", center=(_ILL_PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS) + sim.initialize() + pinball_ids = [3, 4, 5] + + gpu_block(sim, lambda: sim.run(warmup, zero_obs=True)) + log(f" Warmup done ({warmup} steps)") + gpu_block(sim, lambda: sim.snapshot()) + + # ---- Step 3: Stage0 (zero rotation) ---- + log("Step 3: Stage0 (zero rotation)...") + gpu_block(sim, lambda: sim.restore()) + stage0 = run_stage(sim, None, sensor_ids, pinball_ids, target_sensor, + np.zeros(3, dtype=np.float32), dtw_norm_scale, si) + + # ---- Step 4: Stage1 (reference open-loop) ---- + log("Step 4: Stage1 (reference open-loop)...") + gpu_block(sim, lambda: sim.restore()) + stage1 = run_stage(sim, None, sensor_ids, pinball_ids, target_sensor, + _ILL_REF_OMEGA, dtw_norm_scale, si) + + sim.close() + log(" CFD closed.") + + # ---- Step 5: Compute calibration ---- + log("Step 5: Computing calibration constants...") + + all_forces = np.concatenate([stage0["obs_slices"][:, 6:12], + stage1["obs_slices"][:, 6:12]], axis=0) + force_scale = float(np.max(np.abs(all_forces))) + force_scale = max(force_scale, 0.0001) + + all_sens = np.concatenate([stage0["obs_slices"][:, 0:6], + stage1["obs_slices"][:, 0:6]], axis=0) + sens_scale = float(np.max(np.abs(all_sens))) + sens_scale = max(sens_scale, 0.01) + + s0_sim = float(np.mean(stage0["sims"])) + s1_sim = float(np.mean(stage1["sims"])) + + sim_bp = [0.0, s0_sim, s1_sim, + s1_sim + (1.0 - s1_sim) * 0.4, + s1_sim + (1.0 - s1_sim) * 0.7, + 1.0] + sim_val = [0.0, 0.2, 0.5, 0.8, 0.9, 1.0] + + for i in range(1, len(sim_bp)): + if sim_bp[i] <= sim_bp[i - 1]: + sim_bp[i] = sim_bp[i - 1] + 0.01 + + # Rounding + force_scale = round(force_scale, 4) + sens_scale = round(sens_scale, 2) + dtw_norm_scale = round(dtw_norm_scale, 3) + sim_bp = [round(x, 2) for x in sim_bp] + for i in range(1, len(sim_bp)): + if sim_bp[i] <= sim_bp[i - 1]: + sim_bp[i] = sim_bp[i - 1] + 0.01 + sim_bp[-1] = 1.0 + + log(f" FORCE_SCALE = {force_scale:.4f}, SENS_SCALE = {sens_scale:.2f}") + log(f" Stage0 sim = {s0_sim:.4f}, Stage1 sim = {s1_sim:.4f}") + log(f" SIM_BP = {[f'{x:.2f}' for x in sim_bp]}") + + # ---- Step 6: Write calibration.json ---- + calibration = { + "case": case, "scene": "illusion", + "grid": {"nx": NX, "ny": NY}, + "config_path": config_path, + "SI": si, + "FIFO_LEN": FIFO_LEN, "CONV_LEN": CONV_LEN, + "SENSOR_CC": SENSOR_CC, + "FORCE_SCALE": force_scale, "SENS_SCALE": sens_scale, + "dtw_norm_scale": float(dtw_norm_scale), + "SIM_BP": [float(x) for x in sim_bp], + "SIM_VAL": [float(x) for x in sim_val], + "K_CD": K_CD, "K_CL": K_CL, + "W_CD": W_CD, "W_CL": W_CL, "W_SIM": W_SIM, + "FLOOR_CD": FLOOR_CD, "FLOOR_CL": FLOOR_CL, "FLOOR_SIM": FLOOR_SIM, + "FLOOR_PENALTY": FLOOR_PENALTY, + "ACTION_SCALE": ACTION_SCALE, "ACTION_BIAS": ACTION_BIAS.tolist(), + "U0": U0, "RADIUS": RADIUS, "L0": L0, + } + with open(out_dir / "calibration.json", "w") as f: + json.dump(calibration, f, indent=2) + + log(f"\nCalibration complete. Files in {out_dir}:") + log(f" calibration.json, target.npy, target_harmonics.json, calibrate.log") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main() -> int: + parser = argparse.ArgumentParser(description="Phase 0 calibration") + parser.add_argument("--case", type=str, required=True) + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--si", type=int, default=800) + parser.add_argument("--out-dir", type=str, default=None) + parser.add_argument("--scene", type=str, default="karman", + choices=["karman", "illusion"], + help="Scene type (karman or illusion)") + args = parser.parse_args() + + case = args.case + scene = args.scene + config_path = args.config + device_id = args.device_id + si = args.si + warmup = int(4.0 * NX / U0) + + if args.out_dir: + out_dir = Path(args.out_dir).resolve() + else: + out_dir = Path(__file__).resolve().parent / "calibrations" / case + out_dir.mkdir(parents=True, exist_ok=True) + + log_path = out_dir / "calibrate.log" + def log(msg): + line = f"[{time.strftime('%H:%M:%S')}] {msg}" + print(line, flush=True) + with open(log_path, "a") as f: + f.write(line + "\n"); f.flush() + + log(f"=== Phase 0 Calibration: {case} ({scene}) ===") + log(f" Config: {config_path}") + log(f" SI: {si}") + log(f" Output: {out_dir}") + + if scene == "illusion": + _calibrate_illusion(case, config_path, device_id, si, out_dir, log, warmup) + return 0 + + # ---- Step 1: Record target (Karman: dist_cyl + sensors only) ---- + log("Step 1: Recording target signal...") + _clean_cache() + sim = Simulation(lbm_config_path=config_path, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + dist_id_t = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0) + s0_t = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0) + s1_t = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0) + s2_t = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0) + sim.initialize() + + gpu_block(sim, lambda: sim.run(warmup, zero_obs=True)) + log(f" Warmup done ({warmup} steps)") + + target = np.zeros((FIFO_LEN, 6), dtype=np.float32) + for i in range(FIFO_LEN): + gpu_block(sim, lambda: sim.run(si, zero_obs=True)) + target[i] = [ + sim.read_sensor(s0_t, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s0_t, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s1_t, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s1_t, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s2_t, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s2_t, normalize=True)[1] * SENSOR_CC, + ] + sim.close() + np.save(str(out_dir / "target.npy"), target) + target_std = np.std(target, axis=0) + uy_std_avg = float(np.mean([target_std[1], target_std[3], target_std[5]])) + dtw_norm_scale = max(uy_std_avg, 0.01) + log(f" Target recorded. s1_uy std={target[:, 3].std():.4f}") + log(f" DTW norm_scale: {dtw_norm_scale:.4f}") + + # ---- Step 2: Create training sim ---- + log("Step 2: Creating training sim (all 7 objects)...") + _clean_cache() + sim = Simulation(lbm_config_path=config_path, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + dist_id = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0) + sensor_ids = [ + sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0), + sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0), + sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0), + ] + sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS) + sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS) + sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS) + sim.initialize() + pinball_ids = [4, 5, 6] + + gpu_block(sim, lambda: sim.run(warmup, zero_obs=True)) + log(f" Warmup done ({warmup} steps)") + gpu_block(sim, lambda: sim.snapshot()) + + # ---- Step 3: Stage0 (zero rotation) ---- + log("Step 3: Stage0 (zero rotation)...") + gpu_block(sim, lambda: sim.restore()) + stage0 = run_stage(sim, dist_id, sensor_ids, pinball_ids, target, + np.zeros(3, dtype=np.float32), dtw_norm_scale, si) + + # ---- Step 4: Stage1 (reference open-loop) ---- + log("Step 4: Stage1 (reference open-loop)...") + gpu_block(sim, lambda: sim.restore()) + stage1 = run_stage(sim, dist_id, sensor_ids, pinball_ids, target, + _REF_OMEGA, dtw_norm_scale, si) + + sim.close() + log(" CFD closed.") + + # ---- Step 5: Compute calibration ---- + log("Step 5: Computing calibration constants...") + + all_forces = np.concatenate([stage0["obs_slices"][:, 6:12], + stage1["obs_slices"][:, 6:12]], axis=0) + force_scale = float(np.max(np.abs(all_forces))) + force_scale = max(force_scale, 0.0001) + + all_sens = np.concatenate([stage0["obs_slices"][:, 0:6], + stage1["obs_slices"][:, 0:6]], axis=0) + sens_scale = float(np.max(np.abs(all_sens))) + sens_scale = max(sens_scale, 0.01) + + log(f" FORCE_SCALE = {force_scale:.6f}") + log(f" SENS_SCALE = {sens_scale:.4f}") + + s0_sim = float(np.mean(stage0["sims"])) + s1_sim = float(np.mean(stage1["sims"])) + log(f" Stage0 sim_mean = {s0_sim:.4f}") + log(f" Stage1 sim_mean = {s1_sim:.4f}") + + sim_bp = [0.0, s0_sim, s1_sim, + s1_sim + (1.0 - s1_sim) * 0.4, + s1_sim + (1.0 - s1_sim) * 0.7, + 1.0] + sim_val = [0.0, 0.2, 0.5, 0.8, 0.9, 1.0] + + for i in range(1, len(sim_bp)): + if sim_bp[i] <= sim_bp[i - 1]: + sim_bp[i] = sim_bp[i - 1] + 0.01 + + log(f" SIM_BP (raw) = {[f'{x:.4f}' for x in sim_bp]}") + log(f" SIM_VAL = {sim_val}") + + # ---- Step 5b: Round to clean values ---- + # Round measured values to reasonable precision so calibration files + # are human-readable and robust to small measurement noise. + force_scale = round(force_scale, 4) # 0.002429 -> 0.0024 + sens_scale = round(sens_scale, 2) # 0.7543 -> 0.75 + dtw_norm_scale = round(dtw_norm_scale, 3) # 0.2043 -> 0.204 + sim_bp = [round(x, 2) for x in sim_bp] # 0.3166 -> 0.32 etc. + # Ensure sim_bp monotonicity preserved after rounding + for i in range(1, len(sim_bp)): + if sim_bp[i] <= sim_bp[i - 1]: + sim_bp[i] = sim_bp[i - 1] + 0.01 + # Clamp last to 1.0 + sim_bp[-1] = 1.0 + + log(f" FORCE_SCALE (rounded) = {force_scale:.4f}") + log(f" SENS_SCALE (rounded) = {sens_scale:.2f}") + log(f" dtw_norm_scale (rounded) = {dtw_norm_scale:.3f}") + log(f" SIM_BP (rounded) = {[f'{x:.2f}' for x in sim_bp]}") + + # ---- Step 6: Write calibration.json ---- + calibration = { + "case": case, + "grid": {"nx": NX, "ny": NY}, + "config_path": config_path, + "SI": si, + "FIFO_LEN": FIFO_LEN, + "CONV_LEN": CONV_LEN, + "SENSOR_CC": SENSOR_CC, + "FORCE_SCALE": force_scale, + "SENS_SCALE": sens_scale, + "dtw_norm_scale": float(dtw_norm_scale), + "SIM_BP": [float(x) for x in sim_bp], + "SIM_VAL": [float(x) for x in sim_val], + "K_CD": K_CD, + "K_CL": K_CL, + "W_CD": W_CD, + "W_CL": W_CL, + "W_SIM": W_SIM, + "FLOOR_CD": FLOOR_CD, + "FLOOR_CL": FLOOR_CL, + "FLOOR_SIM": FLOOR_SIM, + "FLOOR_PENALTY": FLOOR_PENALTY, + "ACTION_SCALE": ACTION_SCALE, + "ACTION_BIAS": ACTION_BIAS.tolist(), + "U0": U0, + "RADIUS": RADIUS, + "L0": L0, + } + + with open(out_dir / "calibration.json", "w") as f: + json.dump(calibration, f, indent=2) + + log(f"\nCalibration complete. Files in {out_dir}:") + log(f" calibration.json, target.npy, calibrate.log") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/drl_pinball/train/calibrations/illusion_1L/calibration.json b/src/drl_pinball/train/calibrations/illusion_1L/calibration.json new file mode 100644 index 0000000..3dbc1b6 --- /dev/null +++ b/src/drl_pinball/train/calibrations/illusion_1L/calibration.json @@ -0,0 +1,50 @@ +{ + "case": "illusion_1L", + "scene": "illusion", + "grid": { + "nx": 2000, + "ny": 600 + }, + "config_path": "/home/frank14f/DynamisLab/configs/config_lbm_karman_2000x600.json", + "SI": 600, + "FIFO_LEN": 150, + "CONV_LEN": 30, + "SENSOR_CC": 78.0, + "FORCE_SCALE": 0.0029, + "SENS_SCALE": 0.93, + "dtw_norm_scale": 0.251, + "SIM_BP": [ + 0.0, + 0.81, + 0.82, + 0.83, + 0.84, + 1.0 + ], + "SIM_VAL": [ + 0.0, + 0.2, + 0.5, + 0.8, + 0.9, + 1.0 + ], + "K_CD": 50.0, + "K_CL": 100.0, + "W_CD": 0.3, + "W_CL": 0.3, + "W_SIM": 0.4, + "FLOOR_CD": 0.1, + "FLOOR_CL": 0.1, + "FLOOR_SIM": 0.1, + "FLOOR_PENALTY": 0.05, + "ACTION_SCALE": 12.0, + "ACTION_BIAS": [ + 0.0, + 0.0, + 0.0 + ], + "U0": 0.01, + "RADIUS": 10.0, + "L0": 20.0 +} \ No newline at end of file diff --git a/src/drl_pinball/train/calibrations/illusion_1L/target_harmonics.json b/src/drl_pinball/train/calibrations/illusion_1L/target_harmonics.json new file mode 100644 index 0000000..e94d202 --- /dev/null +++ b/src/drl_pinball/train/calibrations/illusion_1L/target_harmonics.json @@ -0,0 +1,194 @@ +[ + { + "dc": 0.6848811427752177, + "amps": [ + 0.20348444790268552, + 0.02915900547523486, + 0.02784031012371259, + 0.026538733180849892, + 0.01919840514104942 + ], + "freqs": [ + 0.02666666666666667, + 0.03333333333333333, + 0.02, + 0.05333333333333334, + 0.04666666666666667 + ], + "phases": [ + -1.2362005505810982, + -1.1766679461536396, + 1.8863477863626301, + -2.834029481822404, + -0.7286212732915559 + ] + }, + { + "dc": -0.026989686206604045, + "amps": [ + 0.2790809917571785, + 0.07572308632544687, + 0.04473095779630388, + 0.039113569175253486, + 0.038411479291876154 + ], + "freqs": [ + 0.02666666666666667, + 0.05333333333333334, + 0.02, + 0.04666666666666667, + 0.03333333333333333 + ], + "phases": [ + 3.0547687538583306, + 0.6711306028525605, + -0.10499736264498984, + -2.7955698201715835, + 3.1220861518480647 + ] + }, + { + "dc": 0.5674182403087616, + "amps": [ + 0.06514963980633084, + 0.02472242867395058, + 0.014768360330077245, + 0.010156019183366295, + 0.00963919337572356 + ], + "freqs": [ + 0.05333333333333334, + 0.04666666666666667, + 0.060000000000000005, + 0.04, + 0.1 + ], + "phases": [ + -0.7488772609862894, + 2.444199252505853, + -0.7923018439616549, + 2.505537406241232, + -1.0758943656055078 + ] + }, + { + "dc": 0.014232361110819814, + "amps": [ + 0.4461623192829029, + 0.0779793315804837, + 0.0550568453231777, + 0.05331603051277121, + 0.04874852247514895 + ], + "freqs": [ + 0.02666666666666667, + 0.02, + 0.08, + 0.03333333333333333, + 0.07333333333333333 + ], + "phases": [ + -3.1049303564692066, + 0.027203636269757633, + 0.09727455160562376, + -3.095222555392279, + -3.044905056768087 + ] + }, + { + "dc": 0.6885532836119334, + "amps": [ + 0.2018316634259717, + 0.033110505502103905, + 0.029481533925939486, + 0.027209600513899444, + 0.0210247548932384 + ], + "freqs": [ + 0.02666666666666667, + 0.05333333333333334, + 0.02, + 0.03333333333333333, + 0.07333333333333333 + ], + "phases": [ + 1.8817739461803038, + 2.742416499482124, + -1.1159813908192473, + 1.7448219790472927, + 1.8563271918371878 + ] + }, + { + "dc": 0.0464287880451108, + "amps": [ + 0.266793500575173, + 0.08445240294008177, + 0.05549481483267607, + 0.03176632041902732, + 0.025647047642157136 + ], + "freqs": [ + 0.02666666666666667, + 0.05333333333333334, + 0.02, + 0.013333333333333334, + 0.04666666666666667 + ], + "phases": [ + 3.0338443035387885, + -2.634676303912741, + -0.03215847385538798, + -0.005548069481591528, + 0.8284453225261742 + ] + }, + { + "dc": 0.002807500216489037, + "amps": [ + 1.712089159739782e-05, + 6.84709236632855e-06, + 3.6652517552615438e-06, + 2.9729489026324893e-06, + 2.0179845909227172e-06 + ], + "freqs": [ + 0.05333333333333334, + 0.04666666666666667, + 0.060000000000000005, + 0.04, + 0.06666666666666667 + ], + "phases": [ + -2.4049945951829694, + 0.6548426535356905, + -2.327050319985697, + 0.5677103875435425, + -2.2688377992989093 + ] + }, + { + "dc": -1.4845058922219323e-05, + "amps": [ + 0.0006606466483213061, + 0.00010502206909109318, + 8.38977596666718e-05, + 4.8485264824562914e-05, + 4.61590173098796e-05 + ], + "freqs": [ + 0.02666666666666667, + 0.02, + 0.03333333333333333, + 0.013333333333333334, + 0.04 + ], + "phases": [ + -0.8350712750146696, + 2.4283127217217664, + -0.9184054668600495, + 2.6043773129885874, + -0.9754872507212389 + ] + } +] \ No newline at end of file diff --git a/src/drl_pinball/train/calibrations/re100/calibration.json b/src/drl_pinball/train/calibrations/re100/calibration.json new file mode 100644 index 0000000..113c2e9 --- /dev/null +++ b/src/drl_pinball/train/calibrations/re100/calibration.json @@ -0,0 +1,49 @@ +{ + "case": "re100", + "grid": { + "nx": 2000, + "ny": 600 + }, + "config_path": "/home/frank14f/DynamisLab/configs/config_lbm_karman_2000x600.json", + "SI": 800, + "FIFO_LEN": 150, + "CONV_LEN": 30, + "SENSOR_CC": 78.0, + "FORCE_SCALE": 0.0024, + "SENS_SCALE": 0.75, + "dtw_norm_scale": 0.204, + "SIM_BP": [ + 0.0, + 0.32, + 0.82, + 0.89, + 0.94, + 1.0 + ], + "SIM_VAL": [ + 0.0, + 0.2, + 0.5, + 0.8, + 0.9, + 1.0 + ], + "K_CD": 50.0, + "K_CL": 100.0, + "W_CD": 0.3, + "W_CL": 0.3, + "W_SIM": 0.4, + "FLOOR_CD": 0.1, + "FLOOR_CL": 0.1, + "FLOOR_SIM": 0.1, + "FLOOR_PENALTY": 0.05, + "ACTION_SCALE": 12.0, + "ACTION_BIAS": [ + 0.0, + 0.0, + 0.0 + ], + "U0": 0.01, + "RADIUS": 10.0, + "L0": 20.0 +} \ No newline at end of file diff --git a/src/drl_pinball/train/env_illusion.py b/src/drl_pinball/train/env_illusion.py new file mode 100644 index 0000000..6461945 --- /dev/null +++ b/src/drl_pinball/train/env_illusion.py @@ -0,0 +1,426 @@ +#!/usr/bin/env python3 +"""Hydrodynamic Illusion environment (V5 - calibration-driven, no_bias, 2000x600). + +Two-phase initialization: + Phase 1: target cylinder (diameter=1.0*L0) + 3 sensors -> warmup -> + record 150-step signal -> FFT harmonics -> close + Phase 2: Training sim (3 sensors + 3 pinball) -> warmup -> + zero-action FIFO -> snapshot + +Observation (14-dim, physical norm, NO clip): + [0:6] = raw_forces / FORCE_SCALE (front_fx,fy, top_fx,fy, bot_fx,fy) + [6:12] = raw_sensors / SENS_SCALE (s0_ux,uy, s1_ux,uy, s2_ux,uy) + [12:14] = target_cd, target_cl (from harmonics reconstruction) + +Action (3-dim): no_bias only + [-1,1] -> omega = -(action * 12 + [0,0,0]) * U0 / R + +Reward: Gaussian cd/cl (compared to harmonics-reconstructed target) + normalized DTW. +""" +from __future__ import annotations + +import json, os, sys, time +from collections import deque +from pathlib import Path +from typing import Optional, Tuple + +import numpy as np +import gymnasium as gym +from gymnasium import spaces + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from CelerisLab import Simulation + +_CELERIS = Path("/home/frank14f/CelerisLab") +_CONFIG_H = _CELERIS / "src/CelerisLab/lbm/kernels/config/config_objects.h" +_PTX = _CELERIS / "src/CelerisLab/lbm/kernels/kernel.ptx" + +def _clean_cache(): + for p in [_CONFIG_H, _PTX]: + if p.exists(): p.unlink() + +# --------------------------------------------------------------------------- +L0 = 20.0; U0 = 0.01; RADIUS = L0 / 2.0 +NX = 2000; NY = 600 +CENTER_Y = float(NY - 1) / 2.0 + +# Illusion geometry: pinball and sensors shifted left vs Karman +PINBALL_FRONT_X = 380.0 # 19*L0 +PINBALL_REAR_X = 406.0 # 20.3*L0 +SENSOR_X = 600.0 # 30*L0 +TARGET_CYL_X = 400.0 # 20*L0 (target cylinder position during recording) +TARGET_CYL_RADIUS = 1.0 * L0 # 1L target diameter + +FIFO_LEN = 150; CONV_LEN = 30; MAX_STEPS = 500 +EMA_FAST = 0.2 +S_DIM = 14; A_DIM = 3 +SENSOR_CC = 78.0 +ACTION_SCALE = 12.0 +ACTION_BIAS = np.array([0.0, 0.0, 0.0], dtype=np.float32) + +# --------------------------------------------------------------------------- +# DTW utilities +# --------------------------------------------------------------------------- +def calc_lag(target, state): + t_mean = np.mean(target); s_mean = np.mean(state) + corr = np.correlate(target - t_mean, state - s_mean, mode="full") + lags = np.arange(-len(target) + 1, len(target)) + return int(lags[np.argmax(corr)]) + +def calc_dtw_sim(target, state, norm_scale=1.0): + n, m = len(target), len(state) + dtw = np.full((n + 1, m + 1), np.inf); dtw[0, 0] = 0.0 + for i in range(1, n + 1): + for j in range(1, m + 1): + cost = abs(float(target[i - 1]) - float(state[j - 1])) + last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1]) + dtw[i, j] = cost + last_min + raw = 1.0 - dtw[n, m] / (float(n) * norm_scale) + return float(max(0.0, raw)) + +def compute_similarity(target, fifo_arr, conv_len, norm_scale): + target = np.asarray(target, dtype=np.float64) + state = np.asarray(fifo_arr, dtype=np.float64) + if len(state) < conv_len: + return 0.0 + id_sens = 3 + target_seq = target[conv_len:2 * conv_len, id_sens] + state_seq = state[-conv_len:, id_sens] + lag = calc_lag(target_seq, state_seq) + sim = 0.0 + for i in range(6): + t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len] + s_seq = state[-conv_len:, i] + sim += calc_dtw_sim(t_seq, s_seq, norm_scale=norm_scale) + return float(sim / 6.0) + + +class ActionSmoother: + def __init__(self, weight=0.1): + self.weight = weight; self._state = None + def __call__(self, target): + t = np.asarray(target, dtype=np.float32) + if self._state is None: + self._state = t.copy() + else: + self._state = (1.0 - self.weight) * self._state + self.weight * t + return self._state.copy() + def reset(self, value=None): + self._state = np.asarray(value, dtype=np.float32).copy() if value is not None else None + + +# --------------------------------------------------------------------------- +# Harmonics analysis +# --------------------------------------------------------------------------- +def analyze_harmonics(states, n_harmonics=5): + N, D = states.shape + result = [] + for d in range(D): + y = states[:, d] + fft_coef = np.fft.rfft(y) + freqs = np.fft.rfftfreq(N, d=1) + amps = 2 * np.abs(fft_coef) / N + phases = np.angle(fft_coef) + idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1 + harmonics = { + 'dc': float(np.real(fft_coef[0]) / N), + 'amps': amps[idx].tolist(), + 'freqs': freqs[idx].tolist(), + 'phases': phases[idx].tolist(), + } + result.append(harmonics) + return result + + +def gen_target_states_at(t, harmonics): + t = np.asarray(t) + D = len(harmonics) + result = np.zeros((t.size, D), dtype=np.float32) + for d, h in enumerate(harmonics): + val = np.full(t.shape, h['dc'], dtype=np.float32) + for amp, freq, phase in zip(h['amps'], h['freqs'], h['phases']): + val += amp * np.cos(2 * np.pi * freq * t + phase) + result[:, d] = val + if result.shape[0] == 1: + return result[0] + return result + + +# --------------------------------------------------------------------------- +def record_illusion_target(config_path, device_id, si, target_diam=1.0): + _clean_cache() + warmup = int(4.0 * NX / U0) + sim = Simulation(lbm_config_path=config_path, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + sim.add_body("circle", center=(TARGET_CYL_X, CENTER_Y, 0.0), + radius=target_diam * L0) + s0 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0) + s1 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0) + s2 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0) + sim.initialize() + sim.run(warmup, zero_obs=True) + + target_states = np.zeros((FIFO_LEN, 8), dtype=np.float32) + for i in range(FIFO_LEN): + sim.run(si, zero_obs=True) + target_states[i] = [ + sim.read_sensor(s0, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s0, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s1, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s1, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s2, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s2, normalize=True)[1] * SENSOR_CC, + sim.read_force(0, normalize=True)[0], + sim.read_force(0, normalize=True)[1], + ] + sim.close() + harmonics = analyze_harmonics(target_states, n_harmonics=5) + return target_states, harmonics + + +# --------------------------------------------------------------------------- +class IllusionCloakEnv(gym.Env): + metadata = {"render_modes": ["human"]} + + def __init__(self, device_id=0, seed=42, calibration=None, + config_path=None, target_states=None, target_harmonics=None, + target_diam=1.0): + super().__init__() + self.device_id = device_id + self.seed = seed + np.random.seed(seed) + self._target_diam = target_diam + + if calibration is None: + raise ValueError("calibration dict is required") + self._cal = calibration.copy() + self._si = int(self._cal["SI"]) + self._force_scale = np.float32(self._cal["FORCE_SCALE"]) + self._sens_scale = np.float32(self._cal["SENS_SCALE"]) + self._dtw_norm_scale = float(self._cal["dtw_norm_scale"]) + self._sim_bp = np.array(self._cal["SIM_BP"], dtype=np.float64) + self._sim_val = np.array(self._cal["SIM_VAL"], dtype=np.float64) + self._k_cd = float(self._cal["K_CD"]); self._k_cl = float(self._cal["K_CL"]) + self._w_cd = float(self._cal["W_CD"]); self._w_cl = float(self._cal["W_CL"]) + self._w_sim = float(self._cal["W_SIM"]) + self._floor_cd = float(self._cal["FLOOR_CD"]) + self._floor_cl = float(self._cal["FLOOR_CL"]) + self._floor_sim = float(self._cal["FLOOR_SIM"]) + self._floor_pen = float(self._cal["FLOOR_PENALTY"]) + + self._config_path = config_path or self._cal.get("config_path") + if self._config_path is None: + raise ValueError("config_path is required") + self._target_states = target_states + self._target_harmonics = target_harmonics + + self.action_space = spaces.Box(-1.0, 1.0, (A_DIM,), dtype=np.float32) + self.observation_space = spaces.Box(-10.0, 10.0, (S_DIM,), dtype=np.float32) + + self.fifo_states = deque(maxlen=FIFO_LEN) + self.save_states = None + self.target_sensor = None # (FIFO_LEN, 6) for DTW + self.current_step = 0 + self.smoother = ActionSmoother(weight=0.1) + self._ema_r_cd = 0.0; self._ema_r_cl = 0.0 + self.sim = None + self.sensor_ids = [] + self.pinball_ids = [] + self._init_cfd() + + def _gpu_block(self, fn): + if self.sim is not None: + self.sim.ctx._ctx.push() + try: + fn() + finally: + if self.sim is not None: + self.sim.ctx._ctx.pop() + + def _init_cfd(self): + t0 = time.perf_counter() + warmup = int(4.0 * NX / U0) + + # ---- Phase 1: Target ---- + if self._target_states is None or self._target_harmonics is None: + print(" [illusion] Phase 1: Recording target + harmonics...", flush=True) + self._target_states, self._target_harmonics = record_illusion_target( + self._config_path, self.device_id, self._si, self._target_diam) + print(" [illusion] Target recorded.", flush=True) + + self.target_sensor = self._target_states[:, 0:6].copy() + + # ---- Phase 2: Training sim ---- + print(" [illusion] Phase 2: Training sim...", flush=True) + _clean_cache() + self.sim = Simulation(lbm_config_path=self._config_path, device_id=self.device_id) + self.sim._assert_object_count_contract = lambda *a, **kw: None + self.sensor_ids = [ + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0), + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0), + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0), + ] + self.sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS) + self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS) + self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS) + self.sim.initialize() + self.pinball_ids = [3, 4, 5] + + self._gpu_block(lambda: self.sim.run(warmup, zero_obs=True)) + print(f" [illusion] Warmup done ({time.perf_counter()-t0:.0f}s).") + + zero_omega = self._action_to_omega(np.zeros(3, dtype=np.float32)) + self.smoother.reset(zero_omega.copy()) + fifo_save = [] + for _ in range(FIFO_LEN): + self._set_omega(zero_omega) + self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True)) + obs = self._read_obs() + sl = obs[0:6] * SENSOR_CC + fifo_save.append(sl.copy()) + self.save_states = np.array(fifo_save, dtype=np.float32) + + self._gpu_block(lambda: self.sim.snapshot()) + print(f" [illusion] Init done ({time.perf_counter()-t0:.0f}s)") + + def _read_obs(self): + obs = [] + for sid in self.sensor_ids: + s = self.sim.read_sensor(sid, normalize=True) + obs.extend([float(s[0]), float(s[1])]) + for pid in self.pinball_ids: + obs.extend(self.sim.read_force(pid, normalize=True)) + return np.array(obs, dtype=np.float32) + + def _action_to_omega(self, action_norm): + sv = (np.asarray(action_norm, dtype=np.float32) * ACTION_SCALE + ACTION_BIAS) * U0 + return -sv / RADIUS + + def _set_omega(self, omega): + for pid, w in zip(self.pinball_ids, omega): + self.sim.set_body(pid, omega=float(w)) + + def _normalize_obs(self, raw_obs): + forces = raw_obs[6:12] / self._force_scale + sens = raw_obs[0:6] / self._sens_scale + return np.hstack([forces, sens]).astype(np.float32) + + def _make_obs(self, obs_slice, target_cd, target_cl): + base = self._normalize_obs(obs_slice) + return np.hstack([base, np.float32(target_cd), np.float32(target_cl)]) + + def _compute_reward(self, obs_slice): + forces_raw = obs_slice[6:12] + cd_raw = forces_raw[0] + forces_raw[2] + forces_raw[4] + cl_raw = forces_raw[1] + forces_raw[3] + forces_raw[5] + cd_norm = cd_raw / self._force_scale + cl_norm = cl_raw / self._force_scale + + tgt = gen_target_states_at(self.current_step, self._target_harmonics) + target_cd = float(tgt[6]) / self._force_scale + target_cl = float(tgt[7]) / self._force_scale + + sim_val = 0.0 + if len(self.fifo_states) >= CONV_LEN * 2: + sim_val = compute_similarity(self.target_sensor, + np.array(list(self.fifo_states)), + CONV_LEN, self._dtw_norm_scale) + + r_cd_raw = float(np.exp(-(cd_norm - target_cd)**2 * self._k_cd)) + r_cl_raw = float(np.exp(-(cl_norm - target_cl)**2 * self._k_cl)) + self._ema_r_cd = (1 - EMA_FAST) * self._ema_r_cd + EMA_FAST * r_cd_raw + self._ema_r_cl = (1 - EMA_FAST) * self._ema_r_cl + EMA_FAST * r_cl_raw + r_sim = float(np.interp(sim_val, self._sim_bp, self._sim_val)) + + reward = self._w_cd * self._ema_r_cd + self._w_cl * self._ema_r_cl + self._w_sim * r_sim + + floor_pen = 0.0 + if self._ema_r_cd < self._floor_cd: + floor_pen += self._floor_pen * (self._floor_cd - self._ema_r_cd) / self._floor_cd + if self._ema_r_cl < self._floor_cl: + floor_pen += self._floor_pen * (self._floor_cl - self._ema_r_cl) / self._floor_cl + if r_sim < self._floor_sim: + floor_pen += self._floor_pen * (self._floor_sim - r_sim) / self._floor_sim + reward = max(0.0, reward - floor_pen) + + info = {"cd": float(cd_norm), "cl": float(cl_norm), + "target_cd": float(target_cd), "target_cl": float(target_cl), + "sim": float(sim_val), + "r_cd": self._ema_r_cd, "r_cl": self._ema_r_cl, "r_sim": r_sim, + "floor_pen": float(floor_pen)} + return float(reward), info + + def reset(self, seed=None, options=None): + super().reset(seed=seed) + self._gpu_block(lambda: self.sim.restore()) + self.smoother.reset(self._action_to_omega(np.zeros(3, dtype=np.float32))) + self.fifo_states.clear() + for i in range(len(self.save_states)): + self.fifo_states.append(self.save_states[i, :]) + self.current_step = 0 + self._ema_r_cd = 0.0; self._ema_r_cl = 0.0 + obs_raw = self._read_obs() + tgt = gen_target_states_at(self.current_step, self._target_harmonics) + tgt_cd = float(tgt[6]) / self._force_scale + tgt_cl = float(tgt[7]) / self._force_scale + obs = self._make_obs(obs_raw, tgt_cd, tgt_cl) + return obs, {} + + def step(self, action): + assert self.action_space.contains(action), f"Invalid action: {action}" + target_omega = self._action_to_omega(action) + smoothed = self.smoother(target_omega) + self._set_omega(smoothed) + + self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True)) + + self.current_step += 1 + obs_raw = self._read_obs() + self.fifo_states.append(obs_raw[0:6] * SENSOR_CC) + reward, info = self._compute_reward(obs_raw) + + tgt = gen_target_states_at(self.current_step, self._target_harmonics) + tgt_cd = float(tgt[6]) / self._force_scale + tgt_cl = float(tgt[7]) / self._force_scale + obs = self._make_obs(obs_raw, tgt_cd, tgt_cl) + + terminated = False + return obs, reward, terminated, False, info + + def render(self, mode="human"): + pass + + def close(self): + if self.sim is not None: + self.sim.close() + + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--calibration", type=str, required=True) + parser.add_argument("--config", type=str, required=True) + args = parser.parse_args() + + with open(args.calibration, "r") as f: + cal = json.load(f) + + print("=== IllusionCloakEnv V5 Quick Test ===") + env = IllusionCloakEnv(device_id=args.device_id, calibration=cal, + config_path=args.config) + obs, _ = env.reset() + print(f" Init obs: min={obs.min():.4f}, max={obs.max():.4f}, mean={obs.mean():.4f}") + rewards = [] + for step in range(50): + obs, reward, *_ = env.step(np.zeros(3, dtype=np.float32)) + rewards.append(reward) + print(f" Zero-action reward (last 20): {np.mean(rewards[-20:]):.4f}") + obs1, _ = env.reset() + obs2, _ = env.reset() + print(f" Reset consistency: {np.max(np.abs(obs1-obs2)):.8f}") + env.close() + print("=== Done ===") diff --git a/src/drl_pinball/train/env_karman.py b/src/drl_pinball/train/env_karman.py new file mode 100644 index 0000000..3bf47f2 --- /dev/null +++ b/src/drl_pinball/train/env_karman.py @@ -0,0 +1,419 @@ +#!/usr/bin/env python3 +"""Karman Cloak environment (V5 — parameterized, no-bias only). + +Based on env_karman_2000x600.py V4. Accepts calibration dict and config path. +All calibration constants loaded from calibration.json produced by calibrate.py. + +Design: Two-phase initialization to AVOID runtime sync_bodies(). + Phase 1: Temporary Simulation(dist + sensors) -> record target -> close + Phase 2: Training Simulation(all objects upfront) -> warmup -> zero-action FIFO -> snapshot + +CUDA context: mirrors legacy pattern — push CFD context before GPU ops, pop after. + +Observation (12-dim, physical norm, NO clip — VecNormalize handles that): + [0:6] = raw_forces / FORCE_SCALE (front_fx,fy, top_fx,fy, bot_fx,fy) + [6:12] = raw_sensors / SENS_SCALE (s0_ux,uy, s1_ux,uy, s2_ux,uy) + +Action (3-dim): no_bias only + [-1,1] -> omega = -(action * ACTION_SCALE + [0,0,0]) * U0 / R + +Reward (V3: Gaussian + EMA smoothing + normalized DTW): + r_cd = EMA(exp(-cd_norm^2 * K_CD), EMA_FAST) + r_cl = EMA(exp(-cl_norm^2 * K_CL), EMA_FAST) + r_sim = piecewise_map(sim, SIM_BP, SIM_VAL) -> [0, 1] + reward = W_CD*r_cd + W_CL*r_cl + W_SIM*r_sim - floor_penalty +""" +from __future__ import annotations + +import json +import os, sys, time +from collections import deque +from pathlib import Path +from typing import Optional, Tuple + +import numpy as np +import gymnasium as gym +from gymnasium import spaces + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from CelerisLab import Simulation + +_CELERIS = Path("/home/frank14f/CelerisLab") +_CONFIG_H = _CELERIS / "src/CelerisLab/lbm/kernels/config/config_objects.h" +_PTX = _CELERIS / "src/CelerisLab/lbm/kernels/kernel.ptx" + +def _clean_cache(): + for p in [_CONFIG_H, _PTX]: + if p.exists(): p.unlink() + +# --------------------------------------------------------------------------- +# Geometry constants (fixed across all Karman cloak cases) +# --------------------------------------------------------------------------- +L0 = 20.0; D_CYL = L0; U0 = 0.01; RADIUS = L0 / 2.0 +NX = 2000; NY = 600 +CENTER_Y = float(NY - 1) / 2.0 + +DIST_X = 600.0 +PINBALL_FRONT_X = 1000.0 +PINBALL_REAR_X = 1026.0 +SENSOR_X = 1200.0 + +FIFO_LEN = 150; CONV_LEN = 30; MAX_STEPS = 500 +EMA_FAST = 0.2 +S_DIM = 12; A_DIM = 3 + +SENSOR_CC = 78.0 + + +# --------------------------------------------------------------------------- +# DTW utilities (identical to env_karman_2000x600.py) +# --------------------------------------------------------------------------- +def calc_lag(target: np.ndarray, state: np.ndarray) -> int: + t_mean = np.mean(target); s_mean = np.mean(state) + corr = np.correlate(target - t_mean, state - s_mean, mode="full") + lags = np.arange(-len(target) + 1, len(target)) + return int(lags[np.argmax(corr)]) + + +def calc_dtw_sim(target: np.ndarray, state: np.ndarray, + norm_scale: float = 1.0) -> float: + n, m = len(target), len(state) + dtw = np.full((n + 1, m + 1), np.inf); dtw[0, 0] = 0.0 + for i in range(1, n + 1): + for j in range(1, m + 1): + cost = abs(float(target[i - 1]) - float(state[j - 1])) + last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1]) + dtw[i, j] = cost + last_min + raw = 1.0 - dtw[n, m] / (float(n) * norm_scale) + return float(max(0.0, raw)) + + +def compute_similarity(target_states, fifo_states, conv_len, norm_scale): + target = np.asarray(target_states, dtype=np.float64) + state = np.asarray(fifo_states, dtype=np.float64) + id_sens = 3 + target_seq = target[conv_len:2 * conv_len, id_sens] + state_seq = state[-conv_len:, id_sens] + lag = calc_lag(target_seq, state_seq) + sim = 0.0 + for i in range(6): + t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len] + s_seq = state[-conv_len:, i] + sim += calc_dtw_sim(t_seq, s_seq, norm_scale=norm_scale) + return float(sim / 6.0) + + +# --------------------------------------------------------------------------- +class ActionSmoother: + def __init__(self, weight: float = 0.1): + self.weight = weight; self._state: Optional[np.ndarray] = None + def __call__(self, target: np.ndarray) -> np.ndarray: + t = np.asarray(target, dtype=np.float32) + if self._state is None: + self._state = t.copy() + else: + self._state = (1.0 - self.weight) * self._state + self.weight * t + return self._state.copy() + def reset(self, value: Optional[np.ndarray] = None) -> None: + self._state = np.asarray(value, dtype=np.float32).copy() if value is not None else None + + +# --------------------------------------------------------------------------- +def record_target(config_path: str, device_id: int, si: int) -> np.ndarray: + """Record target signal (dist_cyl + sensors, no pinball).""" + _clean_cache() + warmup = int(4.0 * NX / U0) + sim = Simulation(lbm_config_path=config_path, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0) + s0 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0) + s1 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0) + s2 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0) + sim.initialize() + sim.run(warmup, zero_obs=True) + target = np.zeros((FIFO_LEN, 6), dtype=np.float32) + for i in range(FIFO_LEN): + sim.run(si, zero_obs=True) + target[i] = [ + sim.read_sensor(s0, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s0, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s1, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s1, normalize=True)[1] * SENSOR_CC, + sim.read_sensor(s2, normalize=True)[0] * SENSOR_CC, + sim.read_sensor(s2, normalize=True)[1] * SENSOR_CC, + ] + sim.close() + return np.array(target, dtype=np.float32) + + +# --------------------------------------------------------------------------- +class KarmanCloakEnv(gym.Env): + """Parameterized Karman Cloak environment (V5 — no_bias only). + + Parameters + ---------- + device_id : int + GPU device ID. + seed : int + Random seed. + calibration : dict + Calibration dict loaded from calibration.json. + Must contain: FORCE_SCALE, SENS_SCALE, dtw_norm_scale, + SIM_BP, SIM_VAL, K_CD, K_CL, W_CD, W_CL, W_SIM, + FLOOR_CD, FLOOR_CL, FLOOR_SIM, FLOOR_PENALTY, + ACTION_SCALE, ACTION_BIAS, SI. + config_path : str + Path to LBM config JSON. + target_states : np.ndarray, optional + Pre-recorded target signal. If None, recorded on-the-fly. + """ + metadata = {"render_modes": ["human"]} + + def __init__(self, device_id: int = 0, seed: int = 42, + calibration: Optional[dict] = None, + config_path: Optional[str] = None, + target_states: Optional[np.ndarray] = None): + super().__init__() + self.device_id = device_id + self.seed = seed + np.random.seed(seed) + + # Load calibration + if calibration is None: + raise ValueError("calibration dict is required for V5 KarmanCloakEnv") + self._cal = calibration.copy() + self._si = int(self._cal["SI"]) + self._force_scale = np.float32(self._cal["FORCE_SCALE"]) + self._sens_scale = np.float32(self._cal["SENS_SCALE"]) + self._dtw_norm_scale = float(self._cal["dtw_norm_scale"]) + self._sim_bp = np.array(self._cal["SIM_BP"], dtype=np.float64) + self._sim_val = np.array(self._cal["SIM_VAL"], dtype=np.float64) + self._k_cd = float(self._cal["K_CD"]) + self._k_cl = float(self._cal["K_CL"]) + self._w_cd = float(self._cal["W_CD"]) + self._w_cl = float(self._cal["W_CL"]) + self._w_sim = float(self._cal["W_SIM"]) + self._floor_cd = float(self._cal["FLOOR_CD"]) + self._floor_cl = float(self._cal["FLOOR_CL"]) + self._floor_sim = float(self._cal["FLOOR_SIM"]) + self._floor_pen = float(self._cal["FLOOR_PENALTY"]) + self._action_scale = float(self._cal["ACTION_SCALE"]) + self._action_bias = np.array(self._cal["ACTION_BIAS"], dtype=np.float32) + + self._config_path = config_path or self._cal.get("config_path") + if self._config_path is None: + raise ValueError("config_path is required") + + self.action_space = spaces.Box(-1.0, 1.0, (A_DIM,), dtype=np.float32) + self.observation_space = spaces.Box(-10.0, 10.0, (S_DIM,), dtype=np.float32) + + self.fifo_states = deque(maxlen=FIFO_LEN) + self.save_states: np.ndarray = None + self.target_states: np.ndarray = None + self.current_step = 0 + self.smoother = ActionSmoother(weight=0.1) + self._ema_r_cd = 0.0 + self._ema_r_cl = 0.0 + self.sim = None + self.dist_id = None + self.sensor_ids = [] + self.pinball_ids = [] + + if target_states is not None: + self.target_states = target_states + self._init_cfd() + + def _gpu_block(self, fn): + if self.sim is not None: + self.sim.ctx._ctx.push() + try: + fn() + finally: + if self.sim is not None: + self.sim.ctx._ctx.pop() + + def _init_cfd(self): + t0 = time.perf_counter() + warmup = int(4.0 * NX / U0) + + if self.target_states is None: + self.target_states = record_target(self._config_path, self.device_id, self._si) + + _clean_cache() + self.sim = Simulation(lbm_config_path=self._config_path, device_id=self.device_id) + self.sim._assert_object_count_contract = lambda *a, **kw: None + + self.dist_id = self.sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0) + self.sensor_ids = [ + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0), + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0), + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0), + ] + self.sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS) + self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS) + self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS) + self.sim.initialize() + self.pinball_ids = [4, 5, 6] + + print(f" [env] Warmup ({warmup} steps)...", end=" ", flush=True) + self._gpu_block(lambda: self.sim.run(warmup, zero_obs=True)) + print(f"done ({time.perf_counter()-t0:.0f}s).") + + # Zero-action FIFO: no rotation, just let pinball oscillate naturally + print(f" [env] Zero-action FIFO ({FIFO_LEN})...", end=" ", flush=True) + zero_omega = self._action_to_omega(np.zeros(3, dtype=np.float32)) + self.smoother.reset(zero_omega.copy()) + self._gpu_block(lambda: [self.sim.run(self._si, zero_obs=True) for _ in range(FIFO_LEN)]) + f_diag = self._read_obs() + print(f"max|force|={np.max(np.abs(f_diag[6:12])):.6f}") + + print(f" [env] Saving snapshot after zero-action FIFO...", end=" ", flush=True) + fifo_save = [] + for _ in range(FIFO_LEN): + self._set_omega(zero_omega) + self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True)) + obs = self._read_obs() + sl = obs[2:14].copy() + sl[0:6] *= SENSOR_CC + fifo_save.append(sl) + self.save_states = np.array(fifo_save, dtype=np.float32) + print("done.") + + self._gpu_block(lambda: self.sim.snapshot()) + print(f" [env] Init done ({time.perf_counter()-t0:.0f}s)") + + def _read_obs(self) -> np.ndarray: + obs = list(self.sim.read_force(self.dist_id, normalize=True)) + for sid in self.sensor_ids: + s = self.sim.read_sensor(sid, normalize=True) + obs.extend([float(s[0]), float(s[1])]) + for pid in self.pinball_ids: + obs.extend(self.sim.read_force(pid, normalize=True)) + return np.array(obs, dtype=np.float32) + + def _action_to_omega(self, action_norm: np.ndarray) -> np.ndarray: + sv = (np.asarray(action_norm, dtype=np.float32) * self._action_scale + self._action_bias) * U0 + return -sv / RADIUS + + def _set_omega(self, omega: np.ndarray): + for pid, w in zip(self.pinball_ids, omega): + self.sim.set_body(pid, omega=float(w)) + + def _normalize_obs(self, raw_obs_slice: np.ndarray) -> np.ndarray: + forces = raw_obs_slice[6:12] / self._force_scale + sens = raw_obs_slice[0:6] / self._sens_scale + return np.hstack([forces, sens]).astype(np.float32) + + def _compute_reward(self, obs_slice: np.ndarray) -> Tuple[float, dict]: + forces_raw = obs_slice[6:12] + cd_raw = (forces_raw[0] + forces_raw[2] + forces_raw[4]) / 3.0 + cl_raw = (forces_raw[1] + forces_raw[3] + forces_raw[5]) / 3.0 + cd_norm = cd_raw / self._force_scale + cl_norm = cl_raw / self._force_scale + + sim_val = 0.0 + if len(self.fifo_states) >= CONV_LEN * 2: + sim_val = compute_similarity(self.target_states, + np.array(list(self.fifo_states)), + conv_len=CONV_LEN, + norm_scale=self._dtw_norm_scale) + + r_cd_raw = float(np.exp(-cd_norm**2 * self._k_cd)) + r_cl_raw = float(np.exp(-cl_norm**2 * self._k_cl)) + + self._ema_r_cd = (1 - EMA_FAST) * self._ema_r_cd + EMA_FAST * r_cd_raw + self._ema_r_cl = (1 - EMA_FAST) * self._ema_r_cl + EMA_FAST * r_cl_raw + r_sim = float(np.interp(sim_val, self._sim_bp, self._sim_val)) + + reward = self._w_cd * self._ema_r_cd + self._w_cl * self._ema_r_cl + self._w_sim * r_sim + + floor_pen = 0.0 + if self._ema_r_cd < self._floor_cd: + floor_pen += self._floor_pen * (self._floor_cd - self._ema_r_cd) / self._floor_cd + if self._ema_r_cl < self._floor_cl: + floor_pen += self._floor_pen * (self._floor_cl - self._ema_r_cl) / self._floor_cl + if r_sim < self._floor_sim: + floor_pen += self._floor_pen * (self._floor_sim - r_sim) / self._floor_sim + reward = max(0.0, reward - floor_pen) + + info = {"cd": float(cd_norm), "cl": float(cl_norm), "sim": float(sim_val), + "r_cd": self._ema_r_cd, "r_cl": self._ema_r_cl, "r_sim": r_sim, + "floor_pen": float(floor_pen)} + return float(reward), info + + def reset(self, seed=None, options=None) -> Tuple[np.ndarray, dict]: + super().reset(seed=seed) + self._gpu_block(lambda: self.sim.restore()) + self.smoother.reset(self._action_to_omega(np.zeros(3, dtype=np.float32))) + self.fifo_states.clear() + for i in range(len(self.save_states)): + self.fifo_states.append(self.save_states[i, 0:6]) + self.current_step = 0 + self._ema_r_cd = 0.0 + self._ema_r_cl = 0.0 + obs_raw = self._read_obs() + obs = self._normalize_obs(obs_raw[2:14]) + return obs, {} + + def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, dict]: + assert self.action_space.contains(action), f"Invalid action: {action}" + target_omega = self._action_to_omega(action) + smoothed = self.smoother(target_omega) + self._set_omega(smoothed) + + self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True)) + + obs_raw = self._read_obs() + obs_slice = obs_raw[2:14] + obs = self._normalize_obs(obs_slice) + self.fifo_states.append(obs_slice[0:6] * SENSOR_CC) + reward, info = self._compute_reward(obs_slice) + + self.current_step += 1 + terminated = False + return obs, reward, terminated, False, info + + def render(self, mode="human"): + pass + + def close(self): + if self.sim is not None: + self.sim.close() + + +# --------------------------------------------------------------------------- +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--calibration", type=str, required=True, + help="Path to calibration.json") + parser.add_argument("--config", type=str, required=True, + help="Path to LBM config JSON") + args = parser.parse_args() + + with open(args.calibration, "r") as f: + cal = json.load(f) + + print("=== KarmanCloakEnv V5 Quick Test ===") + env = KarmanCloakEnv(device_id=args.device_id, calibration=cal, + config_path=args.config) + + obs, _ = env.reset() + print(f" Init obs: min={obs.min():.4f}, max={obs.max():.4f}, mean={obs.mean():.4f}") + + rewards = [] + for step in range(50): + obs, reward, *_ = env.step(np.zeros(3, dtype=np.float32)) + rewards.append(reward) + print(f" Zero-action reward (last 20): {np.mean(rewards[-20:]):.4f}") + + obs1, _ = env.reset() + obs2, _ = env.reset() + print(f" Reset consistency: {np.max(np.abs(obs1-obs2)):.8f}") + + env.close() + print("=== Done ===") diff --git a/src/drl_pinball/train/env_karman_2000x600.py b/src/drl_pinball/train/env_karman_2000x600.py index 24c487e..7391277 100644 --- a/src/drl_pinball/train/env_karman_2000x600.py +++ b/src/drl_pinball/train/env_karman_2000x600.py @@ -15,11 +15,11 @@ Observation (12-dim, physical norm, NO clip — VecNormalize handles that): Action (3-dim): [-1,1] -> omega = -(action*8 + [0,-4,4]) * U0 / R -Reward (V2: Gaussian + EMA smoothing + delayed r_sim): - r_cd = EMA(exp(-cd_norm^2 * 50), 0.2) # Gaussian, no zero-crossing spikes - r_cl = EMA(exp(-cl_norm^2 * 100), 0.2) # Gaussian, smoothed - r_sim = delayed_EMA(piecewise_map(sim), 0.05, skip=20) # slow, stable - reward = 0.25*r_cd + 0.25*r_cl + 0.50*r_sim # sim-heavy (most stable) +Reward (V3: Gaussian + EMA smoothing + normalized DTW): + r_cd = EMA(exp(-cd_norm^2 * K_CD), EMA_FAST) # Gaussian, no zero-crossing spikes + r_cl = EMA(exp(-cl_norm^2 * K_CL), EMA_FAST) # Gaussian, smoothed + r_sim = piecewise_map(sim, SIM_BP, SIM_VAL) # normalized DTW, mapped to [0,1] + reward = W_CD*r_cd + W_CL*r_cl + W_SIM*r_sim - floor_penalty """ from __future__ import annotations @@ -70,7 +70,7 @@ K_CL = 100.0 # Gaussian: exp(-cl_norm^2 * K) EMA_FAST = 0.2 # EMA weight for r_cd/r_cl smoothing (r_sim NOT smoothed) # Normalized DTW sim mapping (based on Phase 0: Stage0~0.36, Stage1~0.82) SIM_BP = np.array([0.0, 0.36, 0.82, 0.90, 0.95, 1.0], dtype=np.float64) -SIM_VAL = np.array([0.0, 0.2, 0.5, 0.8, 0.9, 0.95], dtype=np.float64) +SIM_VAL = np.array([0.0, 0.2, 0.5, 0.8, 0.9, 1.0], dtype=np.float64) W_CD = 0.30 W_CL = 0.30 W_SIM = 0.40 diff --git a/src/drl_pinball/train/env_vortex.py b/src/drl_pinball/train/env_vortex.py new file mode 100644 index 0000000..13caa5c --- /dev/null +++ b/src/drl_pinball/train/env_vortex.py @@ -0,0 +1,378 @@ +#!/usr/bin/env python3 +"""Vortex Cloak environment (V5 — no_bias, 2000x600, transfer from Karman). + +Two-phase initialization: + Phase 1: sensors only -> warmup -> add_vortex(x=10) -> record target(150 steps) + Phase 2: sensors + pinball -> warmup -> add_vortex(x=15) -> FIFO -> snapshot + +reset() restores to vortex+pinball state. MAX_STEPS=150 (time-bounded). +After step 150, done=True. No disturbance cylinder. + +Observation (12-dim, physical norm, NO clip): + [0:6] = raw_forces / FORCE_SCALE (front_fx,fy, top_fx,fy, bot_fx,fy) + [6:12] = raw_sensors / SENS_SCALE (s0_ux,uy, s1_ux,uy, s2_ux,uy) + +Action (3-dim): no_bias only + [-1,1] -> omega = -(action * 12 + [0,0,0]) * U0 / R + +Reward: Gaussian + EMA + normalized DTW, same as Karman cloak. +""" +from __future__ import annotations + +import json +import sys, time +from collections import deque +from pathlib import Path +from typing import Optional, Tuple + +import numpy as np +import gymnasium as gym +from gymnasium import spaces + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from CelerisLab import Simulation +from CelerisLab.lbm.initializers import add_vortex + +_CELERIS = Path("/home/frank14f/CelerisLab") +_CONFIG_H = _CELERIS / "src/CelerisLab/lbm/kernels/config/config_objects.h" +_PTX = _CELERIS / "src/CelerisLab/lbm/kernels/kernel.ptx" + +def _clean_cache(): + for p in [_CONFIG_H, _PTX]: + if p.exists(): p.unlink() + +# --------------------------------------------------------------------------- +# Geometry constants +# --------------------------------------------------------------------------- +L0 = 20.0; U0 = 0.01; RADIUS = L0 / 2.0 +NX = 2000; NY = 600 +CENTER_Y = float(NY - 1) / 2.0 + +PINBALL_FRONT_X = 1000.0 +PINBALL_REAR_X = 1026.0 +SENSOR_X = 1200.0 + +VORTEX_X_TARGET = 200.0 # x=10*L0 for target phase +VORTEX_X_TRAIN = 300.0 # x=15*L0 for training (closer to pinball) +VORTEX_RADIUS = 40.0 # 2*L0 +VORTEX_STRENGTH_LAMB = 0.5 * U0 +VORTEX_STRENGTH_TAYLOR = 0.03 * U0 + +FIFO_LEN = 150; CONV_LEN = 30; MAX_STEPS = 150 +EMA_FAST = 0.2 +S_DIM = 12; A_DIM = 3 +SENSOR_CC = 78.0 +ACTION_SCALE = 12.0 +ACTION_BIAS = np.array([0.0, 0.0, 0.0], dtype=np.float32) + +# --------------------------------------------------------------------------- +# DTW utilities (same as env_karman.py) +# --------------------------------------------------------------------------- +def calc_lag(target, state): + t_mean = np.mean(target); s_mean = np.mean(state) + corr = np.correlate(target - t_mean, state - s_mean, mode="full") + lags = np.arange(-len(target) + 1, len(target)) + return int(lags[np.argmax(corr)]) + +def calc_dtw_sim(target, state, norm_scale=1.0): + n, m = len(target), len(state) + dtw = np.full((n + 1, m + 1), np.inf); dtw[0, 0] = 0.0 + for i in range(1, n + 1): + for j in range(1, m + 1): + cost = abs(float(target[i - 1]) - float(state[j - 1])) + last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1]) + dtw[i, j] = cost + last_min + raw = 1.0 - dtw[n, m] / (float(n) * norm_scale) + return float(max(0.0, raw)) + +def compute_similarity(target, fifo_arr, conv_len, norm_scale): + target = np.asarray(target, dtype=np.float64) + state = np.asarray(fifo_arr, dtype=np.float64) + if len(state) < conv_len: + return 0.0 + t_slice = target[:conv_len] + s_slice = state[-conv_len:] + sim = 0.0 + for i in range(6): + sim += calc_dtw_sim(t_slice[:, i], s_slice[:, i], norm_scale=norm_scale) + return float(sim / 6.0) + + +class ActionSmoother: + def __init__(self, weight=0.1): + self.weight = weight; self._state = None + def __call__(self, target): + t = np.asarray(target, dtype=np.float32) + if self._state is None: + self._state = t.copy() + else: + self._state = (1.0 - self.weight) * self._state + self.weight * t + return self._state.copy() + def reset(self, value=None): + self._state = np.asarray(value, dtype=np.float32).copy() if value is not None else None + + +# --------------------------------------------------------------------------- +class VortexCloakEnv(gym.Env): + metadata = {"render_modes": ["human"]} + + def __init__(self, device_id=0, seed=42, calibration=None, + config_path=None, vortex_type="lamb"): + super().__init__() + self.device_id = device_id + self.seed = seed + np.random.seed(seed) + self._vortex_type = vortex_type + + if calibration is None: + raise ValueError("calibration dict is required") + self._cal = calibration.copy() + self._si = int(self._cal["SI"]) + self._force_scale = np.float32(self._cal["FORCE_SCALE"]) + self._sens_scale = np.float32(self._cal["SENS_SCALE"]) + self._dtw_norm_scale = float(self._cal["dtw_norm_scale"]) + self._sim_bp = np.array(self._cal["SIM_BP"], dtype=np.float64) + self._sim_val = np.array(self._cal["SIM_VAL"], dtype=np.float64) + self._k_cd = float(self._cal["K_CD"]); self._k_cl = float(self._cal["K_CL"]) + self._w_cd = float(self._cal["W_CD"]); self._w_cl = float(self._cal["W_CL"]) + self._w_sim = float(self._cal["W_SIM"]) + self._floor_cd = float(self._cal["FLOOR_CD"]) + self._floor_cl = float(self._cal["FLOOR_CL"]) + self._floor_sim = float(self._cal["FLOOR_SIM"]) + self._floor_pen = float(self._cal["FLOOR_PENALTY"]) + + self._config_path = config_path or self._cal.get("config_path") + if self._config_path is None: + raise ValueError("config_path is required") + + self.action_space = spaces.Box(-1.0, 1.0, (A_DIM,), dtype=np.float32) + self.observation_space = spaces.Box(-10.0, 10.0, (S_DIM,), dtype=np.float32) + + self.fifo_states = deque(maxlen=FIFO_LEN) + self.save_states = None + self.target_states = None + self.current_step = 0 + self.smoother = ActionSmoother(weight=0.1) + self._ema_r_cd = 0.0; self._ema_r_cl = 0.0 + self.sim = None + self.sensor_ids = [] + self.pinball_ids = [] + self._init_cfd() + + def _gpu_block(self, fn): + if self.sim is not None: + self.sim.ctx._ctx.push() + try: + fn() + finally: + if self.sim is not None: + self.sim.ctx._ctx.pop() + + def _init_cfd(self): + t0 = time.perf_counter() + warmup = int(4.0 * NX / U0) + + # ---- Phase 1: Target (sensors + vortex only, no pinball) ---- + print(" [vortex] Phase 1: Target recording...", flush=True) + _clean_cache() + sim_t = Simulation(lbm_config_path=self._config_path, device_id=self.device_id) + sim_t._assert_object_count_contract = lambda *a, **kw: None + s0 = sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0) + s1 = sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0) + s2 = sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0) + sensor_ids_t = [s0, s1, s2] + sim_t.initialize() + sim_t.run(warmup, zero_obs=True) + print(f" [vortex] Target warmup done ({warmup} steps)") + + # Inject vortex and record + add_vortex(sim_t.field, (VORTEX_X_TARGET, CENTER_Y), + VORTEX_RADIUS, VORTEX_STRENGTH_TAYLOR if self._vortex_type == "taylor" else VORTEX_STRENGTH_LAMB, + self._vortex_type) + + target = np.zeros((MAX_STEPS, 6), dtype=np.float32) + for i in range(MAX_STEPS): + sim_t.run(self._si, zero_obs=True) + target[i] = [ + sim_t.read_sensor(s0, normalize=True)[0] * SENSOR_CC, + sim_t.read_sensor(s0, normalize=True)[1] * SENSOR_CC, + sim_t.read_sensor(s1, normalize=True)[0] * SENSOR_CC, + sim_t.read_sensor(s1, normalize=True)[1] * SENSOR_CC, + sim_t.read_sensor(s2, normalize=True)[0] * SENSOR_CC, + sim_t.read_sensor(s2, normalize=True)[1] * SENSOR_CC, + ] + sim_t.close() + self.target_states = target + + # ---- Phase 2: Training sim (sensors + pinball + vortex) ---- + print(" [vortex] Phase 2: Training sim...", flush=True) + _clean_cache() + self.sim = Simulation(lbm_config_path=self._config_path, device_id=self.device_id) + self.sim._assert_object_count_contract = lambda *a, **kw: None + self.sensor_ids = [ + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0), + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0), + self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0), + ] + self.sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS) + self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS) + self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS) + self.sim.initialize() + self.pinball_ids = [3, 4, 5] + + self._gpu_block(lambda: self.sim.run(warmup, zero_obs=True)) + print(f" [vortex] Pinball warmup done ({warmup} steps)") + + # Inject vortex at training position (closer to pinball) + self._gpu_block(lambda: add_vortex(self.sim.field, (VORTEX_X_TRAIN, CENTER_Y), + VORTEX_RADIUS, + VORTEX_STRENGTH_TAYLOR if self._vortex_type == "taylor" else VORTEX_STRENGTH_LAMB, + self._vortex_type)) + print(f" [vortex] Vortex ({self._vortex_type}) injected at x={VORTEX_X_TRAIN}") + + # Zero-action FIFO (no rotation, vortex evolves with pinball) + zero_omega = self._action_to_omega(np.zeros(3, dtype=np.float32)) + self.smoother.reset(zero_omega.copy()) + fifo_save = [] + for _ in range(FIFO_LEN): + self._set_omega(zero_omega) + self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True)) + obs = self._read_obs() + sl = obs[0:6] # 6 sensor channels (legacy-equiv) + sl = sl * SENSOR_CC + fifo_save.append(sl.copy()) + self.save_states = np.array(fifo_save, dtype=np.float32) + + self._gpu_block(lambda: self.sim.snapshot()) + print(f" [vortex] Init done ({time.perf_counter()-t0:.0f}s)") + + def _read_obs(self): + obs = [] + for sid in self.sensor_ids: + s = self.sim.read_sensor(sid, normalize=True) + obs.extend([float(s[0]), float(s[1])]) + for pid in self.pinball_ids: + obs.extend(self.sim.read_force(pid, normalize=True)) + return np.array(obs, dtype=np.float32) + + def _action_to_omega(self, action_norm): + sv = (np.asarray(action_norm, dtype=np.float32) * ACTION_SCALE + ACTION_BIAS) * U0 + return -sv / RADIUS + + def _set_omega(self, omega): + for pid, w in zip(self.pinball_ids, omega): + self.sim.set_body(pid, omega=float(w)) + + def _normalize_obs(self, raw_obs): + forces = raw_obs[6:12] / self._force_scale + sens = raw_obs[0:6] / self._sens_scale + return np.hstack([forces, sens]).astype(np.float32) + + def _compute_reward(self, obs_slice): + forces_raw = obs_slice[6:12] + cd_raw = (forces_raw[0] + forces_raw[2] + forces_raw[4]) / 3.0 + cl_raw = (forces_raw[1] + forces_raw[3] + forces_raw[5]) / 3.0 + cd_norm = cd_raw / self._force_scale + cl_norm = cl_raw / self._force_scale + + sim_val = 0.0 + if len(self.fifo_states) >= CONV_LEN: + sim_val = compute_similarity(self.target_states, + np.array(list(self.fifo_states)), + CONV_LEN, self._dtw_norm_scale) + + r_cd_raw = float(np.exp(-cd_norm**2 * self._k_cd)) + r_cl_raw = float(np.exp(-cl_norm**2 * self._k_cl)) + self._ema_r_cd = (1 - EMA_FAST) * self._ema_r_cd + EMA_FAST * r_cd_raw + self._ema_r_cl = (1 - EMA_FAST) * self._ema_r_cl + EMA_FAST * r_cl_raw + r_sim = float(np.interp(sim_val, self._sim_bp, self._sim_val)) + + reward = self._w_cd * self._ema_r_cd + self._w_cl * self._ema_r_cl + self._w_sim * r_sim + + floor_pen = 0.0 + if self._ema_r_cd < self._floor_cd: + floor_pen += self._floor_pen * (self._floor_cd - self._ema_r_cd) / self._floor_cd + if self._ema_r_cl < self._floor_cl: + floor_pen += self._floor_pen * (self._floor_cl - self._ema_r_cl) / self._floor_cl + if r_sim < self._floor_sim: + floor_pen += self._floor_pen * (self._floor_sim - r_sim) / self._floor_sim + reward = max(0.0, reward - floor_pen) + + info = {"cd": float(cd_norm), "cl": float(cl_norm), "sim": float(sim_val), + "r_cd": self._ema_r_cd, "r_cl": self._ema_r_cl, "r_sim": r_sim, + "floor_pen": float(floor_pen)} + return float(reward), info + + def reset(self, seed=None, options=None): + super().reset(seed=seed) + self._gpu_block(lambda: self.sim.restore()) + self.smoother.reset(self._action_to_omega(np.zeros(3, dtype=np.float32))) + self.fifo_states.clear() + for i in range(len(self.save_states)): + self.fifo_states.append(self.save_states[i, :]) + self.current_step = 0 + self._ema_r_cd = 0.0; self._ema_r_cl = 0.0 + obs_raw = self._read_obs() + obs = self._normalize_obs(obs_raw) + return obs, {} + + def step(self, action): + assert self.action_space.contains(action), f"Invalid action: {action}" + target_omega = self._action_to_omega(action) + smoothed = self.smoother(target_omega) + self._set_omega(smoothed) + + self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True)) + + obs_raw = self._read_obs() + obs = self._normalize_obs(obs_raw) + self.fifo_states.append(obs_raw[0:6] * SENSOR_CC) + reward, info = self._compute_reward(obs_raw) + + self.current_step += 1 + done = self.current_step >= MAX_STEPS + return obs, reward, done, False, info + + def render(self, mode="human"): + pass + + def close(self): + if self.sim is not None: + self.sim.close() + + +# --------------------------------------------------------------------------- +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--calibration", type=str, required=True) + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--vortex-type", type=str, default="lamb", + choices=["lamb", "taylor"]) + args = parser.parse_args() + + with open(args.calibration, "r") as f: + cal = json.load(f) + + print(f"=== VortexCloakEnv V5 Quick Test ({args.vortex_type}) ===") + env = VortexCloakEnv(device_id=args.device_id, calibration=cal, + config_path=args.config, vortex_type=args.vortex_type) + + obs, _ = env.reset() + print(f" Init obs: min={obs.min():.4f}, max={obs.max():.4f}, mean={obs.mean():.4f}") + + rewards = [] + for step in range(MAX_STEPS): + obs, reward, done, *_ = env.step(np.zeros(3, dtype=np.float32)) + rewards.append(reward) + if done: + break + print(f" Zero-action avg reward: {np.mean(rewards):.4f}") + print(f" Steps: {len(rewards)}") + + env.close() + print("=== Done ===") diff --git a/src/drl_pinball/train/launch_multi.sh b/src/drl_pinball/train/launch_multi.sh new file mode 100755 index 0000000..c96ccad --- /dev/null +++ b/src/drl_pinball/train/launch_multi.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +# launch_multi.sh — Sequential multi-GPU training launcher for server deployment. +# +# Starts Karman Cloak training on multiple GPUs sequentially, with a configurable +# delay between launches to avoid CelerisLab kernel compilation race conditions. +# +# Usage: +# bash launch_multi.sh --case-name re100_karman --seeds 42,43,44,45,46,47 \ +# --gpus 0,1,2,3,4,5 --episodes 500 \ +# --config configs/config_lbm_karman_2000x600.json \ +# --calibration calibrations/re100/calibration.json +# +# # Transfer learning from a base model +# bash launch_multi.sh --case-name re200_karman --seeds 42,43,44 \ +# --gpus 0,1,2 --episodes 500 \ +# --config configs/config_lbm_karman_2000x600_re200.json \ +# --calibration calibrations/re200/calibration.json \ +# --transfer output/re100_karman_seed42/models/best_model.zip +# +# Requirements: +# - conda env pycuda_3_10 +# - CelerisLab at /home/frank14f/CelerisLab +# - train_karman.py, env_karman.py, symmetry_wrapper.py in same directory +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CACHE_H="${HOME}/CelerisLab/src/CelerisLab/lbm/kernels/config/config_objects.h" +CACHE_PTX="${HOME}/CelerisLab/src/CelerisLab/lbm/kernels/kernel.ptx" + +# --- Defaults --- +CASE_NAME="" +SEEDS="" +GPUS="" +EPISODES=500 +CONFIG="" +CALIBRATION="" +TRANSFER="" +DELAY_SECONDS=420 # 7 minutes between launches +CONDA_ENV="pycuda_3_10" + +usage() { + echo "Usage: $0 --case-name NAME --seeds S1,S2,... --gpus G1,G2,... [options]" + echo "" + echo "Required:" + echo " --case-name NAME Case name for output dirs" + echo " --seeds S1,S2 Comma-separated seed values" + echo " --gpus G1,G2 Comma-separated GPU device IDs" + echo " --config PATH LBM config JSON" + echo " --calibration PATH calibration.json" + echo "" + echo "Options:" + echo " --episodes N Total episodes (default: 500)" + echo " --transfer PATH .zip model for transfer learning" + echo " --delay SEC Seconds between launches (default: 420)" + echo " --env NAME Conda env name (default: pycuda_3_10)" + exit 1 +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --case-name) CASE_NAME="$2"; shift 2 ;; + --seeds) SEEDS="$2"; shift 2 ;; + --gpus) GPUS="$2"; shift 2 ;; + --episodes) EPISODES="$2"; shift 2 ;; + --config) CONFIG="$2"; shift 2 ;; + --calibration) CALIBRATION="$2"; shift 2 ;; + --transfer) TRANSFER="$2"; shift 2 ;; + --delay) DELAY_SECONDS="$2"; shift 2 ;; + --env) CONDA_ENV="$2"; shift 2 ;; + *) echo "Unknown option: $1"; usage ;; + esac +done + +# Validate +if [[ -z "$CASE_NAME" || -z "$SEEDS" || -z "$GPUS" || -z "$CONFIG" || -z "$CALIBRATION" ]]; then + echo "ERROR: Missing required arguments." + usage +fi + +IFS=',' read -ra SEED_ARR <<< "$SEEDS" +IFS=',' read -ra GPU_ARR <<< "$GPUS" + +if [[ ${#SEED_ARR[@]} -ne ${#GPU_ARR[@]} ]]; then + echo "ERROR: Number of seeds (${#SEED_ARR[@]}) must match number of GPUs (${#GPU_ARR[@]})." + exit 1 +fi + +echo "=== Multi-GPU Training Launcher ===" +echo " Case: $CASE_NAME" +echo " Seeds: $SEEDS" +echo " GPUs: $GPUS" +echo " Episodes: $EPISODES" +echo " Config: $CONFIG" +echo " Calibration: $CALIBRATION" +echo " Transfer: ${TRANSFER:-none}" +echo " Delay: ${DELAY_SECONDS}s between launches" +echo " Jobs: ${#SEED_ARR[@]}" +echo "" + +# Clean stale cache before starting +rm -f "$CACHE_H" "$CACHE_PTX" +echo " Cleaned kernel cache." + +# Build transfer arg +TRANSFER_ARG="" +if [[ -n "$TRANSFER" ]]; then + TRANSFER_ARG="--transfer-model $TRANSFER" +fi + +mkdir -p "$SCRIPT_DIR/output" + +for i in "${!SEED_ARR[@]}"; do + seed="${SEED_ARR[$i]}" + gpu="${GPU_ARR[$i]}" + run_name="${CASE_NAME}_seed${seed}" + logfile="$SCRIPT_DIR/output/${run_name}/nohup.log" + + mkdir -p "$SCRIPT_DIR/output/${run_name}" + + echo "[$(date '+%H:%M:%S')] Launching seed=$seed on GPU=$gpu ..." + nohup conda run --no-capture-output -n "$CONDA_ENV" python -u \ + "$SCRIPT_DIR/train_karman.py" \ + --case-name "$CASE_NAME" \ + --device-id "$gpu" \ + --seed "$seed" \ + --total-episodes "$EPISODES" \ + --config "$CONFIG" \ + --calibration "$CALIBRATION" \ + $TRANSFER_ARG \ + > "$logfile" 2>&1 & + + echo " PID: $!" + echo " Log: $logfile" + + if [[ $i -lt $((${#SEED_ARR[@]} - 1)) ]]; then + echo " Waiting ${DELAY_SECONDS}s before next launch..." + sleep "$DELAY_SECONDS" + fi +done + +echo "" +echo "All jobs launched. Monitor with:" +echo " tail -f $SCRIPT_DIR/output/${CASE_NAME}_seed*/nohup.log" +echo " tensorboard --logdir $SCRIPT_DIR/output/${CASE_NAME}_seed*/tb" +echo "" +echo "To stop all:" +echo " ps aux | grep train_karman | grep -v grep | awk '{print \$2}' | xargs kill" diff --git a/src/drl_pinball/train/train_illusion.py b/src/drl_pinball/train/train_illusion.py new file mode 100644 index 0000000..85ce767 --- /dev/null +++ b/src/drl_pinball/train/train_illusion.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Train Hydrodynamic Illusion (V5) - parameterized, no_bias, GPU-native. + +Usage: + conda run -n pycuda_3_10 python -u train_illusion.py \ + --case-name illusion_1L --device-id 0 --seed 42 \ + --config configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/illusion_1L/calibration.json \ + --total-episodes 500 +""" +from __future__ import annotations + +import argparse, json, os, sys, time +from pathlib import Path + +import numpy as np +import pycuda.driver as cuda; cuda.init() + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +import torch +from torch.nn import Module as TorchModule +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize +from torch.utils.tensorboard import SummaryWriter + +from env_illusion import IllusionCloakEnv, record_illusion_target +from symmetry_wrapper import SymmetryAugmentWrapper + + +class Sin(TorchModule): + def __init__(self): super().__init__() + def forward(self, x): return torch.sin(x) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Train Illusion V5") + parser.add_argument("--case-name", type=str, required=True) + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--total-episodes", type=int, default=500) + parser.add_argument("--n-steps", type=int, default=2048) + parser.add_argument("--learn-timesteps", type=int, default=2048) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--n-epochs", type=int, default=10) + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--calibration", type=str, required=True) + parser.add_argument("--si", type=int, default=None) + parser.add_argument("--symmetry-prob", type=float, default=0.5) + parser.add_argument("--transfer-model", type=str, default=None) + args = parser.parse_args() + + with open(args.calibration, "r") as f: + cal = json.load(f) + if args.si is not None: + cal["SI"] = int(args.si) + + run_name = f"{args.case_name}_seed{args.seed}" + out_dir = Path(__file__).resolve().parent / "output" / run_name + (out_dir / "models").mkdir(parents=True, exist_ok=True) + + log_path = out_dir / "train.log" + def log(msg): + line = f"[{time.strftime('%H:%M:%S')}] {msg}" + print(line, flush=True) + with open(log_path, "a") as f: f.write(line + "\n"); f.flush() + + writer = SummaryWriter(log_dir=str(out_dir / "tb")) + log(f"=== V5 Illusion {run_name} ===") + log(f" Device={args.device_id}, Seed={args.seed}, Ep={args.total_episodes}") + log(f" Config={args.config}, SI={cal['SI']}") + + cal_copy = dict(cal) + cal_copy["config_path"] = args.config + with open(out_dir / "calibration.json", "w") as f: + json.dump(cal_copy, f, indent=2) + + # Load pre-recorded target + cal_dir = Path(args.calibration).resolve().parent + target_npy = cal_dir / "target.npy" + harmonics_path = cal_dir / "target_harmonics.json" + if target_npy.exists() and harmonics_path.exists(): + target_states = np.load(str(target_npy)) + with open(harmonics_path, "r") as f: + target_harmonics = json.load(f) + log(f" Loaded target from {cal_dir}") + else: + log(" Recording target...") + t0 = time.perf_counter() + config_path = args.config + si_val = cal["SI"] + target_states, target_harmonics = record_illusion_target( + config_path, args.device_id, si_val) + np.save(str(out_dir / "target.npy"), target_states) + with open(out_dir / "target_harmonics.json", "w") as f: + json.dump(target_harmonics, f, indent=2) + log(f" Target recorded in {time.perf_counter()-t0:.0f}s") + + config_path = args.config + t0 = time.perf_counter() + log(" Creating env...") + env = IllusionCloakEnv(device_id=args.device_id, seed=args.seed, + calibration=cal, config_path=config_path, + target_states=target_states, + target_harmonics=target_harmonics) + env = SymmetryAugmentWrapper(env, prob=args.symmetry_prob, seed=args.seed, + rollout_len=args.n_steps) + vec_env = DummyVecEnv([lambda: env]) + vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=False, + clip_obs=10.0, gamma=0.99) + log(f" Env ready in {time.perf_counter()-t0:.0f}s") + + device = torch.device(f"cuda:{args.device_id}") + log(f" PPO device: {device}") + + if args.transfer_model: + log(f" Transfer from: {args.transfer_model}") + model = PPO.load(args.transfer_model, env=vec_env, device=device, + custom_objects={"activation_fn": Sin}) + log(" Loaded base model.") + else: + model = PPO( + "MlpPolicy", + policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]}, + env=vec_env, device=device, + n_steps=args.n_steps, batch_size=args.batch_size, + n_epochs=args.n_epochs, learning_rate=args.lr, + gamma=0.995, verbose=0, + ) + log(" Created from scratch.") + + best_reward = -float("inf") + t_last = time.perf_counter() + norm_path = str(out_dir / "vec_normalize.pkl") + + for ep in range(1, args.total_episodes + 1): + model.learn(total_timesteps=args.learn_timesteps, reset_num_timesteps=False) + + eval_env = model.get_env() + try: + inner = eval_env.venv.envs[0] + if hasattr(inner, 'prob'): + inner.prob = 0.0 + except Exception: + pass + eval_obs = eval_env.reset() + ep_rewards, ep_r_cd, ep_r_cl, ep_r_sim = [], [], [], [] + for _ in range(360): + action, _ = model.predict(eval_obs) + eval_obs, reward, done, info = eval_env.step(action) + ep_rewards.append(float(reward[0])) + inf = info[0] if isinstance(info, list) else info + if "r_cd" in inf: + ep_r_cd.append(float(inf["r_cd"])) + ep_r_cl.append(float(inf["r_cl"])) + ep_r_sim.append(float(inf["r_sim"])) + if done[0]: break + try: + inner = eval_env.venv.envs[0] + if hasattr(inner, 'prob'): + inner.prob = args.symmetry_prob + except Exception: + pass + avg_r = np.mean(ep_rewards[-180:]) if len(ep_rewards) >= 180 else np.mean(ep_rewards) + dt = time.perf_counter() - t_last; t_last = time.perf_counter() + + writer.add_scalar("eval/avg_reward", avg_r, ep) + if ep_r_cd: + writer.add_scalar("eval/r_cd", float(np.mean(ep_r_cd[-180:])), ep) + writer.add_scalar("eval/r_cl", float(np.mean(ep_r_cl[-180:])), ep) + writer.add_scalar("eval/r_sim", float(np.mean(ep_r_sim[-180:])), ep) + + if avg_r > best_reward: + best_reward = avg_r + model.save(str(out_dir / "models" / "best_model.zip")) + vec_env.save(norm_path) + log(f" Ep {ep:3d}: reward={avg_r:.4f} (BEST) r_cd={np.mean(ep_r_cd[-180:]):.3f} " + f"r_cl={np.mean(ep_r_cl[-180:]):.3f} r_sim={np.mean(ep_r_sim[-180:]):.3f}") + elif ep % 5 == 0: + log(f" Ep {ep:3d}: reward={avg_r:.4f} (best={best_reward:.4f}, {dt:.0f}s/ep)") + + if ep % 10 == 0: + model.save(str(out_dir / "models" / f"chkpt_ep{ep}.zip")) + vec_env.save(norm_path) + + model.save(str(out_dir / "models" / "final_model.zip")) + vec_env.save(norm_path) + + meta = {"case_name": args.case_name, "seed": args.seed, + "total_episodes": args.total_episodes, + "best_reward": float(best_reward), + "n_steps": args.n_steps, "batch_size": args.batch_size, + "n_epochs": args.n_epochs, "lr": args.lr, + "config_path": args.config, "calibration_path": args.calibration} + with (out_dir / "meta.json").open("w") as f: + json.dump(meta, f, indent=2) + + env.close() + log(f"Done. Best reward: {best_reward:.4f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/drl_pinball/train/train_karman.py b/src/drl_pinball/train/train_karman.py new file mode 100644 index 0000000..5b3d79f --- /dev/null +++ b/src/drl_pinball/train/train_karman.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Train Karman Cloak (V5) — parameterized, no-bias only, GPU-native. + +CUDA context management: + - PPO model on GPU (PyTorch owns the default context) + - CFD env wraps its calls in sim.ctx._ctx.push()/pop() (PyCUDA context) + +Usage: + conda run -n pycuda_3_10 python -u train_karman.py \ + --case-name re100_karman --device-id 0 --seed 42 \ + --config configs/config_lbm_karman_2000x600.json \ + --calibration calibrations/re100/calibration.json \ + --total-episodes 500 +""" +from __future__ import annotations + +import argparse, json, os, sys, time +from pathlib import Path + +import numpy as np +import pycuda.driver as cuda; cuda.init() + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +import torch +from torch.nn import Module as TorchModule +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize +from torch.utils.tensorboard import SummaryWriter + +from env_karman import KarmanCloakEnv, record_target +from symmetry_wrapper import SymmetryAugmentWrapper + + +class Sin(TorchModule): + def __init__(self): super().__init__() + def forward(self, x): return torch.sin(x) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Train Karman Cloak V5") + parser.add_argument("--case-name", type=str, required=True, + help="Case name for output dir (e.g. re100_karman)") + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--total-episodes", type=int, default=500) + parser.add_argument("--n-steps", type=int, default=2048) + parser.add_argument("--learn-timesteps", type=int, default=2048) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--n-epochs", type=int, default=10) + parser.add_argument("--config", type=str, required=True, + help="Path to LBM config JSON") + parser.add_argument("--calibration", type=str, required=True, + help="Path to calibration.json") + parser.add_argument("--si", type=int, default=None, + help="Override sample interval (default: from calibration)") + parser.add_argument("--symmetry-prob", type=float, default=0.5, + help="G-symmetry augmentation probability (0=off, 0.5=half)") + parser.add_argument("--transfer-model", type=str, default=None, + help="Path to .zip model for transfer learning") + args = parser.parse_args() + + # Load calibration + with open(args.calibration, "r") as f: + cal = json.load(f) + + # Override SI if specified + if args.si is not None: + cal["SI"] = int(args.si) + + config_path = args.config + si_val = cal["SI"] + + run_name = f"{args.case_name}_seed{args.seed}" + out_dir = Path(__file__).resolve().parent / "output" / run_name + (out_dir / "models").mkdir(parents=True, exist_ok=True) + + log_path = out_dir / "train.log" + def log(msg): + line = f"[{time.strftime('%H:%M:%S')}] {msg}" + print(line, flush=True) + with open(log_path, "a") as f: f.write(line + "\n"); f.flush() + + writer = SummaryWriter(log_dir=str(out_dir / "tb")) + log(f"=== V5 {run_name} ===") + log(f" Device={args.device_id}, Seed={args.seed}, Episodes={args.total_episodes}") + log(f" Config={config_path}, SI={si_val}") + log(f" FORCE_SCALE={cal['FORCE_SCALE']:.6f}, SENS_SCALE={cal['SENS_SCALE']:.4f}") + log(f" dtw_norm_scale={cal['dtw_norm_scale']:.4f}") + + # Save calibration and config alongside training output + cal_copy = dict(cal) + cal_copy["config_path"] = config_path + if args.transfer_model: + cal_copy["transfer_model"] = args.transfer_model + with open(out_dir / "calibration.json", "w") as f: + json.dump(cal_copy, f, indent=2) + + # Load target + cal_dir = Path(args.calibration).resolve().parent + target_path = cal_dir / "target.npy" + if target_path.exists(): + target_states = np.load(str(target_path)) + log(f" Loaded target from {target_path}") + else: + log(" Recording target...") + t0 = time.perf_counter() + target_states = record_target(config_path, args.device_id, si_val) + np.save(str(out_dir / "target.npy"), target_states) + log(f" Target recorded in {time.perf_counter()-t0:.0f}s") + + log(" Creating env...") + t0 = time.perf_counter() + env = KarmanCloakEnv(device_id=args.device_id, seed=args.seed, + calibration=cal, config_path=config_path, + target_states=target_states) + env = SymmetryAugmentWrapper(env, prob=args.symmetry_prob, seed=args.seed, + rollout_len=args.n_steps) + vec_env = DummyVecEnv([lambda: env]) + vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=False, + clip_obs=10.0, gamma=0.99) + log(f" Env ready in {time.perf_counter()-t0:.0f}s") + + device = torch.device(f"cuda:{args.device_id}") + log(f" PPO device: {device}") + + if args.transfer_model: + log(f" Transfer learning from: {args.transfer_model}") + model = PPO.load(args.transfer_model, env=vec_env, device=device, + custom_objects={"activation_fn": Sin}) + log(" Loaded base model.") + else: + model = PPO( + "MlpPolicy", + policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]}, + env=vec_env, + device=device, + n_steps=args.n_steps, + batch_size=args.batch_size, + n_epochs=args.n_epochs, + learning_rate=args.lr, + gamma=0.995, + verbose=0, + ) + log(" Created from scratch.") + + best_reward = -float("inf") + t_last = time.perf_counter() + norm_path = str(out_dir / "vec_normalize.pkl") + + for ep in range(1, args.total_episodes + 1): + model.learn(total_timesteps=args.learn_timesteps, reset_num_timesteps=False) + + # Evaluation (disable symmetry for clean policy eval) + eval_env = model.get_env() + try: + inner = eval_env.venv.envs[0] + if hasattr(inner, 'prob'): + inner.prob = 0.0 + except Exception: + pass + eval_obs = eval_env.reset() + ep_rewards = [] + ep_r_cd, ep_r_cl, ep_r_sim = [], [], [] + for _ in range(360): + action, _ = model.predict(eval_obs) + eval_obs, reward, done, info = eval_env.step(action) + ep_rewards.append(float(reward[0])) + inf = info[0] if isinstance(info, list) else info + if "r_cd" in inf: + ep_r_cd.append(float(inf["r_cd"])) + ep_r_cl.append(float(inf["r_cl"])) + ep_r_sim.append(float(inf["r_sim"])) + if done[0]: break + try: + inner = eval_env.venv.envs[0] + if hasattr(inner, 'prob'): + inner.prob = args.symmetry_prob + except Exception: + pass + avg_r = np.mean(ep_rewards[-180:]) if len(ep_rewards) >= 180 else np.mean(ep_rewards) + dt = time.perf_counter() - t_last; t_last = time.perf_counter() + + writer.add_scalar("eval/avg_reward", avg_r, ep) + if ep_r_cd: + writer.add_scalar("eval/r_cd", float(np.mean(ep_r_cd[-180:])), ep) + writer.add_scalar("eval/r_cl", float(np.mean(ep_r_cl[-180:])), ep) + writer.add_scalar("eval/r_sim", float(np.mean(ep_r_sim[-180:])), ep) + + if avg_r > best_reward: + best_reward = avg_r + model.save(str(out_dir / "models" / "best_model.zip")) + vec_env.save(norm_path) + log(f" Ep {ep:3d}: reward={avg_r:.4f} (BEST, {dt:.0f}s/ep) ** " + f"r_cd={np.mean(ep_r_cd[-180:]):.3f} r_cl={np.mean(ep_r_cl[-180:]):.3f} " + f"r_sim={np.mean(ep_r_sim[-180:]):.3f}") + elif ep % 5 == 0: + log(f" Ep {ep:3d}: reward={avg_r:.4f} (best={best_reward:.4f}, {dt:.0f}s/ep)") + + if ep % 10 == 0: + model.save(str(out_dir / "models" / f"chkpt_ep{ep}.zip")) + vec_env.save(norm_path) + + model.save(str(out_dir / "models" / "final_model.zip")) + vec_env.save(norm_path) + + meta = {"case_name": args.case_name, "seed": args.seed, + "total_episodes": args.total_episodes, + "best_reward": float(best_reward), + "n_steps": args.n_steps, "batch_size": args.batch_size, + "n_epochs": args.n_epochs, "lr": args.lr, + "config_path": config_path, "calibration_path": args.calibration, + "transfer_model": args.transfer_model} + with (out_dir / "meta.json").open("w") as f: + json.dump(meta, f, indent=2) + + env.close() + log(f"Done. Best reward: {best_reward:.4f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/drl_pinball/train/visualize_and_analyze.py b/src/drl_pinball/train/visualize_and_analyze.py new file mode 100644 index 0000000..10cbef1 --- /dev/null +++ b/src/drl_pinball/train/visualize_and_analyze.py @@ -0,0 +1,492 @@ +#!/usr/bin/env python3 +"""Post-training flow-field visualization and quantitative analysis. + +Compares control modes for Karman Cloak 2000x600: + - target: disturbance-only reference (no pinball) + - zero: pinball with zero rotation + - bias: open-loop bias action (zeros -> [0,-4,4]*U0) + - bias_drl: trained Bias best model + - nobias_drl: trained NoBias best model + +Outputs to output/flow_analysis_v4/: + vorticity_*.png, macro_*.npz, metrics summary, comparison plots. + +Usage: + conda run -n pycuda_3_10 python -u visualize_and_analyze.py --device-id 0 +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Callable, Dict, List, Optional, Tuple + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pycuda.driver as cuda + +cuda.init() + +_REPO = Path(__file__).resolve().parents[3] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +import torch +from torch.nn import Module as TorchModule +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize +from CelerisLab.common.render import compute_vorticity, render_vorticity_field + +from env_karman_2000x600 import ( + KarmanCloakEnv, + SI, + FIFO_LEN, + CONV_LEN, + NX, + NY, + L0, + U0, + RADIUS, + CENTER_Y, + DIST_X, + PINBALL_FRONT_X, + PINBALL_REAR_X, + SENSOR_X, + SENSOR_CC, + FORCE_SCALE, + compute_similarity, + record_target, + WARMUP_STEPS, + CFG_PATH, +) + +TRAIN_DIR = Path(__file__).resolve().parent +BIAS_RUN = TRAIN_DIR / "output" / "bias_seed42_s2048_e10_v4" +NOBIAS_RUN = TRAIN_DIR / "output" / "nobias_seed42_s2048_e10_v4" + +CYLINDERS_FULL = [ + ((DIST_X, CENTER_Y), 1.0 * L0), + ((PINBALL_FRONT_X, CENTER_Y), RADIUS), + ((PINBALL_REAR_X, CENTER_Y + 15.0), RADIUS), + ((PINBALL_REAR_X, CENTER_Y - 15.0), RADIUS), +] +CYLINDERS_DIST = [((DIST_X, CENTER_Y), 1.0 * L0)] + + +class Sin(TorchModule): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.sin(x) + + +def log(msg: str) -> None: + print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) + + +def save_flow(sim, out_dir: Path, name: str, cylinders, nx: int = NX, ny: int = NY) -> None: + macro = sim.get_macroscopic() + np.savez_compressed( + out_dir / f"macro_{name}.npz", + rho=macro["rho"], ux=macro["ux"], uy=macro["uy"], + ) + vort = compute_vorticity(macro["ux"], macro["uy"]) + render_vorticity_field( + vort, nx=nx, ny=ny, + out_path=str(out_dir / f"vorticity_{name}.png"), + cylinders=cylinders, + ) + + +def save_target_field(device_id: int, out_dir: Path) -> None: + """Disturbance-only reference flow (no pinball).""" + from env_karman_2000x600 import _clean_cache + from CelerisLab import Simulation + + log("Recording target reference flow field...") + _clean_cache() + sim = Simulation(lbm_config_path=CFG_PATH, device_id=device_id) + sim._assert_object_count_contract = lambda *a, **kw: None + sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0) + sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0) + sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0) + sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0) + sim.initialize() + sim.run(WARMUP_STEPS + FIFO_LEN * SI, zero_obs=True) + save_flow(sim, out_dir, "target", CYLINDERS_DIST) + sim.close() + + +def run_manual_rollout( + env: KarmanCloakEnv, + n_steps: int, + *, + action: Optional[np.ndarray] = None, + fixed_omega: Optional[np.ndarray] = None, + action_provider: Optional[Callable[[np.ndarray], np.ndarray]] = None, +) -> Dict[str, np.ndarray]: + """Run rollout bypassing gym step when fixed_omega is set.""" + obs, _ = env.reset() + if fixed_omega is not None: + env.smoother.reset(np.asarray(fixed_omega, dtype=np.float32)) + rewards, r_cd, r_cl, r_sim = [], [], [], [] + cd_vals, cl_vals, sim_vals = [], [], [] + actions = [] + sensor_hist = [] + force_hist = [] + + for _ in range(n_steps): + if fixed_omega is not None: + omega = np.asarray(fixed_omega, dtype=np.float32) + act = np.zeros(3, dtype=np.float32) + elif action_provider is not None: + act = np.asarray(action_provider(obs), dtype=np.float32).flatten() + omega = env.smoother(env._action_to_omega(act)) + else: + act = np.zeros(3, dtype=np.float32) if action is None else np.asarray(action, dtype=np.float32) + omega = env.smoother(env._action_to_omega(act)) + + env._set_omega(omega) + env._gpu_block(lambda: env.sim.run(SI, zero_obs=True)) + + obs_raw = env._read_obs() + obs_slice = obs_raw[2:14] + obs = env._normalize_obs(obs_slice) + env.fifo_states.append(obs_slice[0:6] * SENSOR_CC) + reward, info = env._compute_reward(obs_slice) + + rewards.append(reward) + r_cd.append(info["r_cd"]) + r_cl.append(info["r_cl"]) + r_sim.append(info["r_sim"]) + cd_vals.append(info["cd"]) + cl_vals.append(info["cl"]) + sim_vals.append(info["sim"]) + actions.append(act.copy()) + sensor_hist.append(obs_slice[0:6] * SENSOR_CC) + force_hist.append(obs_slice[6:12].copy()) + + return { + "rewards": np.array(rewards, dtype=np.float64), + "r_cd": np.array(r_cd, dtype=np.float64), + "r_cl": np.array(r_cl, dtype=np.float64), + "r_sim": np.array(r_sim, dtype=np.float64), + "cd": np.array(cd_vals, dtype=np.float64), + "cl": np.array(cl_vals, dtype=np.float64), + "sim": np.array(sim_vals, dtype=np.float64), + "actions": np.array(actions, dtype=np.float32), + "sensors": np.array(sensor_hist, dtype=np.float32), + "forces": np.array(force_hist, dtype=np.float32), + } + + +def run_model_rollout( + env: KarmanCloakEnv, + model_path: Path, + norm_path: Path, + device_id: int, + n_steps: int, + deterministic: bool = True, +) -> Dict[str, np.ndarray]: + vec_env = DummyVecEnv([lambda: env]) + vec_env = VecNormalize.load(str(norm_path), vec_env) + vec_env.training = False + vec_env.norm_reward = False + + model = PPO.load( + str(model_path), + env=vec_env, + device=torch.device(f"cuda:{device_id}"), + custom_objects={"policy_kwargs": {"activation_fn": Sin, "net_arch": [64, 64]}}, + ) + + obs = vec_env.reset() + rewards, r_cd, r_cl, r_sim = [], [], [], [] + cd_vals, cl_vals, sim_vals = [], [], [] + actions = [] + sensor_hist = [] + force_hist = [] + + for _ in range(n_steps): + action, _ = model.predict(obs, deterministic=deterministic) + obs, reward, done, info = vec_env.step(action) + inf = info[0] if isinstance(info, list) else info + rewards.append(float(reward[0])) + r_cd.append(float(inf.get("r_cd", 0.0))) + r_cl.append(float(inf.get("r_cl", 0.0))) + r_sim.append(float(inf.get("r_sim", 0.0))) + cd_vals.append(float(inf.get("cd", 0.0))) + cl_vals.append(float(inf.get("cl", 0.0))) + sim_vals.append(float(inf.get("sim", 0.0))) + actions.append(np.asarray(action[0], dtype=np.float32)) + raw = env._read_obs()[2:14] + sensor_hist.append(raw[0:6] * SENSOR_CC) + force_hist.append(raw[6:12].copy()) + + return { + "rewards": np.array(rewards, dtype=np.float64), + "r_cd": np.array(r_cd, dtype=np.float64), + "r_cl": np.array(r_cl, dtype=np.float64), + "r_sim": np.array(r_sim, dtype=np.float64), + "cd": np.array(cd_vals, dtype=np.float64), + "cl": np.array(cl_vals, dtype=np.float64), + "sim": np.array(sim_vals, dtype=np.float64), + "actions": np.array(actions, dtype=np.float32), + "sensors": np.array(sensor_hist, dtype=np.float32), + "forces": np.array(force_hist, dtype=np.float32), + } + + +def summarize_rollout(name: str, data: Dict[str, np.ndarray], tail: int = 180) -> Dict[str, float]: + sl = slice(-tail, None) if len(data["rewards"]) >= tail else slice(None) + return { + "name": name, + "reward": float(np.mean(data["rewards"][sl])), + "r_cd": float(np.mean(data["r_cd"][sl])), + "r_cl": float(np.mean(data["r_cl"][sl])), + "r_sim": float(np.mean(data["r_sim"][sl])), + "cd_norm": float(np.mean(data["cd"][sl])), + "cl_norm": float(np.mean(data["cl"][sl])), + "sim_raw": float(np.mean(data["sim"][sl])), + "cd_force": float(np.mean((data["forces"][sl, 0] + data["forces"][sl, 2] + data["forces"][sl, 4]) / 3.0)), + "cl_force": float(np.mean((data["forces"][sl, 1] + data["forces"][sl, 3] + data["forces"][sl, 5]) / 3.0)), + } + + +def save_flow_from_env(env: KarmanCloakEnv, out_dir: Path, name: str) -> None: + env.sim.ctx._ctx.push() + try: + save_flow(env.sim, out_dir, name, CYLINDERS_FULL) + finally: + env.sim.ctx._ctx.pop() + + +def plot_vorticity_panel(out_dir: Path, names: List[str], titles: List[str]) -> None: + n = len(names) + fig, axes = plt.subplots(2, 3, figsize=(18, 7)) + axes = axes.flatten() + for ax, name, title in zip(axes, names, titles): + img_path = out_dir / f"vorticity_{name}.png" + if img_path.exists(): + ax.imshow(plt.imread(img_path)) + ax.set_title(title, fontsize=11) + ax.axis("off") + for ax in axes[n:]: + ax.axis("off") + fig.suptitle("Karman Cloak — Vorticity Comparison (2000×600)", fontsize=14) + fig.tight_layout() + fig.savefig(out_dir / "vorticity_panel.png", dpi=150, bbox_inches="tight") + plt.close(fig) + + +def plot_metrics(summary: List[Dict], out_dir: Path) -> None: + names = [s["name"] for s in summary] + x = np.arange(len(names)) + width = 0.2 + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + + for i, key in enumerate(["reward", "r_cd", "r_cl", "r_sim"]): + axes[0].bar(x + (i - 1.5) * width, [s[key] for s in summary], width, label=key) + axes[0].set_xticks(x) + axes[0].set_xticklabels(names, rotation=15) + axes[0].set_ylim(0, 1.05) + axes[0].set_ylabel("Reward component") + axes[0].set_title("Eval reward (last 180 steps mean)") + axes[0].legend() + axes[0].grid(axis="y", alpha=0.3) + + axes[1].bar(x - width / 2, [s["cd_norm"] for s in summary], width, label="|Cd| norm") + axes[1].bar(x + width / 2, [s["cl_norm"] for s in summary], width, label="|Cl| norm") + axes[1].set_xticks(x) + axes[1].set_xticklabels(names, rotation=15) + axes[1].set_ylabel("Force norm") + axes[1].set_title("Hydrodynamic forces (normalized)") + axes[1].legend() + axes[1].grid(axis="y", alpha=0.3) + + fig.tight_layout() + fig.savefig(out_dir / "metrics_comparison.png", dpi=150) + plt.close(fig) + + +def plot_sensor_signals( + target_states: np.ndarray, + rollouts: Dict[str, Dict[str, np.ndarray]], + out_dir: Path, +) -> None: + """Plot center-sensor uy vs target reference.""" + fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True) + + t_ref = target_states[CONV_LEN:2 * CONV_LEN, 3] # center uy + axes[0].plot(t_ref, "k-", lw=2, label="target (disturbance only)") + colors = {"zero": "#d62728", "bias": "#ff7f0e", "bias_drl": "#2ca02c", "nobias_drl": "#1f77b4"} + for name, data in rollouts.items(): + if name == "target": + continue + s = data["sensors"][-CONV_LEN:, 3] + axes[0].plot(s, color=colors.get(name, "gray"), alpha=0.85, label=name) + + axes[0].set_ylabel("Center sensor uy (legacy-equiv)") + axes[0].set_title("Downstream velocity signal vs target") + axes[0].legend(loc="upper right") + axes[0].grid(alpha=0.3) + + for name, data in rollouts.items(): + if name in ("target", "zero"): + continue + axes[1].plot(data["actions"][:, 0], alpha=0.7, label=f"{name} a_front") + axes[1].plot(data["actions"][:, 1], alpha=0.7, ls="--", label=f"{name} a_top") + axes[1].plot(data["actions"][:, 2], alpha=0.7, ls=":", label=f"{name} a_bot") + axes[1].set_xlabel("Step") + axes[1].set_ylabel("Action [-1, 1]") + axes[1].set_title("Control actions") + axes[1].legend(ncol=2, fontsize=8) + axes[1].grid(alpha=0.3) + + fig.tight_layout() + fig.savefig(out_dir / "sensor_action_timeseries.png", dpi=150) + plt.close(fig) + + +def plot_reward_traces(rollouts: Dict[str, Dict[str, np.ndarray]], out_dir: Path) -> None: + fig, ax = plt.subplots(figsize=(12, 4)) + colors = {"zero": "#d62728", "bias": "#ff7f0e", "bias_drl": "#2ca02c", "nobias_drl": "#1f77b4"} + for name, data in rollouts.items(): + if name == "target": + continue + r = data["rewards"] + w = min(30, len(r)) + smooth = np.convolve(r, np.ones(w) / w, mode="same") + ax.plot(smooth, color=colors.get(name, "gray"), label=name, alpha=0.9) + ax.set_xlabel("Step") + ax.set_ylabel("Reward (smoothed)") + ax.set_title("Reward traces during eval rollout") + ax.legend() + ax.grid(alpha=0.3) + fig.tight_layout() + fig.savefig(out_dir / "reward_traces.png", dpi=150) + plt.close(fig) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--target", type=str, default=str(TRAIN_DIR / "target.npy")) + parser.add_argument("--n-steps", type=int, default=360) + parser.add_argument("--out", type=str, default=str(TRAIN_DIR / "output" / "flow_analysis_v4")) + args = parser.parse_args() + + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + target_states = np.load(args.target) + n_steps = args.n_steps + + log(f"Output -> {out_dir}") + log(f"Eval steps per case: {n_steps}") + + # --- Target reference field --- + save_target_field(args.device_id, out_dir) + + rollouts: Dict[str, Dict[str, np.ndarray]] = {} + summary: List[Dict] = [] + + # --- Bias env cases (zero, bias, bias_drl) --- + log("Creating Bias env...") + bias_env = KarmanCloakEnv( + device_id=args.device_id, seed=42, target_states=target_states, + ) + + dtw_norm_scale = bias_env._dtw_norm_scale + + log("Case: zero rotation...") + rollouts["zero"] = run_manual_rollout( + bias_env, n_steps, fixed_omega=np.zeros(3, dtype=np.float32), + ) + save_flow_from_env(bias_env, out_dir, "zero") + + log("Case: bias (open-loop)...") + rollouts["bias"] = run_manual_rollout( + bias_env, n_steps, action=np.zeros(3, dtype=np.float32), + ) + save_flow_from_env(bias_env, out_dir, "bias") + + log("Case: Bias DRL best...") + rollouts["bias_drl"] = run_model_rollout( + bias_env, + BIAS_RUN / "models" / "best_model.zip", + BIAS_RUN / "vec_normalize.pkl", + args.device_id, + n_steps, + deterministic=True, + ) + save_flow_from_env(bias_env, out_dir, "bias_drl") + + bias_env.close() + + # --- NoBias DRL --- + log("Creating NoBias env...") + nobias_env = KarmanCloakEnv( + device_id=args.device_id, seed=42, target_states=target_states, + action_bias=np.array([0.0, 0.0, 0.0], dtype=np.float32), + action_scale=12.0, + ) + + log("Case: NoBias DRL best...") + rollouts["nobias_drl"] = run_model_rollout( + nobias_env, + NOBIAS_RUN / "models" / "best_model.zip", + NOBIAS_RUN / "vec_normalize.pkl", + args.device_id, + n_steps, + deterministic=True, + ) + save_flow_from_env(nobias_env, out_dir, "nobias_drl") + nobias_env.close() + + # --- Summaries --- + for name, data in rollouts.items(): + s = summarize_rollout(name, data) + # DTW on full fifo at end + if len(data["sensors"]) >= CONV_LEN: + fifo = np.zeros((FIFO_LEN, 6), dtype=np.float32) + n = min(FIFO_LEN, len(data["sensors"])) + fifo[-n:] = data["sensors"][-n:] + s["sim_dtw_end"] = float(compute_similarity( + target_states, fifo, conv_len=CONV_LEN, + norm_scale=dtw_norm_scale, + )) + summary.append(s) + log(f" {name:12s}: reward={s['reward']:.4f} r_cd={s['r_cd']:.3f} " + f"r_cl={s['r_cl']:.3f} r_sim={s['r_sim']:.3f} sim_raw={s['sim_raw']:.3f}") + + with open(out_dir / "summary.json", "w") as f: + json.dump(summary, f, indent=2) + + # --- Plots --- + plot_vorticity_panel( + out_dir, + ["target", "zero", "bias", "bias_drl", "nobias_drl"], + ["Target (no pinball)", "Zero rotation", "Bias open-loop", + "Bias DRL (best)", "NoBias DRL (best)"], + ) + plot_metrics(summary, out_dir) + plot_sensor_signals(target_states, rollouts, out_dir) + plot_reward_traces(rollouts, out_dir) + + np.savez_compressed( + out_dir / "rollout_data.npz", + **{f"{k}_{fld}": v for k, d in rollouts.items() for fld, v in d.items()}, + target_states=target_states, + ) + + log(f"Done. Results in {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())