#!/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 # --------------------------------------------------------------------------- # 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) 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) 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 ===")