5f061bec06
- Add crossre_transfer.sh: calibrate → transfer-train for re60→re200→re400 - Add re60 config (ν=0.006667, SI=800, uniform+free-slip, very weak shedding) - Calibrate re60, re200, re400: FORCE_SCALE, SENS_SCALE, dtw_norm_scale, SIM_BP - Fix all paths: use DynamisLab submodule CelerisLab, remove external ~/CelerisLab - Remove _clean_cache() from envs/calibrate — CelerisLab handles internally - Move V4 backups to old/: env_karman_2000x600, train_karman_2000x600, etc. - train_karman.py: save model + vecnormalize every episode (non-optional) - Update TRAIN_KNOWLEDGE.md: file structure, calibration table, cross-re guide - All 3 Re verified: 5-episode transfer test passed (re60: 0.64, re200: 0.43, re400: 0.49) Co-authored-by: Cursor <cursoragent@cursor.com>
369 lines
15 KiB
Python
369 lines
15 KiB
Python
#!/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 ===")
|