b3ee72e144
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>
420 lines
17 KiB
Python
420 lines
17 KiB
Python
#!/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 ===")
|