feat(train): V5 parameterized training pipeline — Karman + Illusion verified
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 <cursoragent@cursor.com>
This commit is contained in:
Frank14f
parent
b1ac9ceabb
commit
b3ee72e144
50
configs/config_lbm_karman_2000x600_re200.json
Normal file
50
configs/config_lbm_karman_2000x600_re200.json
Normal file
@ -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"
|
||||
}
|
||||
}
|
||||
50
configs/config_lbm_karman_2000x600_re400.json
Normal file
50
configs/config_lbm_karman_2000x600_re400.json
Normal file
@ -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"
|
||||
}
|
||||
}
|
||||
50
configs/config_lbm_karman_2000x600_re50.json
Normal file
50
configs/config_lbm_karman_2000x600_re50.json
Normal file
@ -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"
|
||||
}
|
||||
}
|
||||
189
src/drl_pinball/train/SERVER_DEPLOY.md
Normal file
189
src/drl_pinball/train/SERVER_DEPLOY.md
Normal file
@ -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 <repo-url>
|
||||
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 |
|
||||
@ -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.
|
||||
|
||||
573
src/drl_pinball/train/calibrate.py
Normal file
573
src/drl_pinball/train/calibrate.py
Normal file
@ -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())
|
||||
@ -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
|
||||
}
|
||||
@ -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
|
||||
]
|
||||
}
|
||||
]
|
||||
49
src/drl_pinball/train/calibrations/re100/calibration.json
Normal file
49
src/drl_pinball/train/calibrations/re100/calibration.json
Normal file
@ -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
|
||||
}
|
||||
426
src/drl_pinball/train/env_illusion.py
Normal file
426
src/drl_pinball/train/env_illusion.py
Normal file
@ -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 ===")
|
||||
419
src/drl_pinball/train/env_karman.py
Normal file
419
src/drl_pinball/train/env_karman.py
Normal file
@ -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 ===")
|
||||
@ -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
|
||||
|
||||
378
src/drl_pinball/train/env_vortex.py
Normal file
378
src/drl_pinball/train/env_vortex.py
Normal file
@ -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 ===")
|
||||
147
src/drl_pinball/train/launch_multi.sh
Executable file
147
src/drl_pinball/train/launch_multi.sh
Executable file
@ -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"
|
||||
207
src/drl_pinball/train/train_illusion.py
Normal file
207
src/drl_pinball/train/train_illusion.py
Normal file
@ -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())
|
||||
226
src/drl_pinball/train/train_karman.py
Normal file
226
src/drl_pinball/train/train_karman.py
Normal file
@ -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())
|
||||
492
src/drl_pinball/train/visualize_and_analyze.py
Normal file
492
src/drl_pinball/train/visualize_and_analyze.py
Normal file
@ -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())
|
||||
Loading…
Reference in New Issue
Block a user