"""Illusion DRL inference (all S_DIM=14, regardless of model name). All illusion models use 14-D observation space (sensors(6) + forces(6) + target_cd(1) + target_cl(1)), with target forces reconstructed from harmonics. Usage: conda run -n pycuda_3_10 python scripts/collect_illusion.py --device 2 --steps 500 Output: data/illusion/{scene_name}/ """ from __future__ import annotations import argparse import json import os import sys import time from collections import deque import numpy as np _REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) if _REPO not in sys.path: sys.path.insert(0, _REPO) _SRC = os.path.join(_REPO, "src") if _SRC not in sys.path: sys.path.insert(0, _SRC) from LegacyCelerisLab import FlowField from CCD_analysis.configs import get_scene, get_scene_list, data_dir_for_scene, model_path_for_scene, LEGACY_CFG_DIR from CCD_analysis.utils.cfd_interface import ( load_legacy_configs, save_vorticity_png, vorticity_from_ddf, load_ppo_model, scale_action, get_velocity_field, calc_lag, calc_dtw_sim, ) from CCD_analysis.utils.resampling import analyze_harmonics, gen_target_states_at DATA_TYPE = np.float32 L0 = 20.0 CENTER_Y = (512 - 1) / 2.0 FIFO_LEN = 150 CONV_LEN = 36 def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: cfg = get_scene(scene_name) out_dir = data_dir_for_scene(scene_name) u0 = cfg["u0"] si = cfg["sample_interval"] ac_scale = cfg["action_scale"] ac_bias = cfg["action_bias"] n_obj = cfg["n_objects_env"] s_dim = cfg["s_dim"] cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR) field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(u0)) with open(os.path.join(out_dir, "config.json"), "w") as f: json.dump({k: str(v) if not isinstance(v, (int, float, list, bool)) else v for k, v in cfg.items()}, f, indent=2) # === Target recording (separate FlowField) === print("=== Target recording ===") ff_tgt = FlowField(field_cfg, cuda_cfg, device_id=device_id) tgt_radius = cfg["target_diameter"] * L0 ff_tgt.add_cylinder((20.0 * L0, CENTER_Y, 0.0), tgt_radius) print(f" target cylinder: diameter={cfg['target_diameter']}L, radius={tgt_radius}", flush=True) for y_off in [2.0, 0.0, -2.0]: ff_tgt.add_sensor((30.0 * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0) n_tgt = 4 ff_tgt.run(int(4 * 1280 / u0), np.zeros(n_tgt, dtype=DATA_TYPE)) target_states = np.empty((0, 8), dtype=DATA_TYPE) for _ in range(FIFO_LEN): ff_tgt.run(si, np.zeros(n_tgt, dtype=DATA_TYPE)) target_states = np.vstack((target_states, ff_tgt.obs.copy()[0:8])) target_harmonics = analyze_harmonics(target_states, n_harmonics=5) np.savez(os.path.join(out_dir, "target.npz"), target_states=target_states) harm_save = [{k: v for k, v in h.items()} for h in target_harmonics] with open(os.path.join(out_dir, "target_harmonics.json"), "w") as f: json.dump(harm_save, f, indent=2) save_vorticity_png(os.path.join(out_dir, "vorticity_target.png"), vorticity_from_ddf(ff_tgt, u0=u0), title="Illusion target cylinder") del ff_tgt # === Control env (6 objects) === print("=== Pinball env + norm ===") ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) for y_off in [2.0, 0.0, -2.0]: ff.add_sensor((30.0 * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0) ff.add_cylinder((19.0 * L0, CENTER_Y, 0.0), L0 / 2.0) ff.add_cylinder((20.3 * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0) ff.add_cylinder((20.3 * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0) n_env = 6 ff.run(int(4 * 1280 / u0), np.zeros(n_env, dtype=DATA_TYPE)) ff.get_ddf() ff.save_ddf() # Norm fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): ff.run(si, np.zeros(n_env, dtype=DATA_TYPE)) fifo.append(ff.obs.copy()[0:12]) temp = np.array(fifo, dtype=DATA_TYPE) force_norm_fact = 6.0 * float(np.max(np.abs(temp[:, 6:12]))) sens_deviation = np.mean(temp[:, 0:6], axis=0).astype(DATA_TYPE) sens_norm_fact = np.zeros(6, dtype=DATA_TYPE) for i in range(6): sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp[:, i] - sens_deviation[i]))) norm = {"force_norm_fact": force_norm_fact, "sens_deviation": sens_deviation.tolist(), "sens_norm_fact": sens_norm_fact.tolist()} with open(os.path.join(out_dir, "norm.json"), "w") as f: json.dump(norm, f, indent=2) print(f" force_norm_fact={force_norm_fact:.6f}") # Preset-action FIFO init (matches legacy_env_imit: [0,0,0,0,-1*U0,1*U0]) # NOTE: this is NOT the same as action_bias([0,-2,2]). action_bias controls DRL # action scaling; preset_action is a fixed Omega array used to warm up the FIFO. ff.apply_ddf() bias = np.zeros(n_env, dtype=DATA_TYPE) bias[4] = -1.0 * u0 bias[5] = 1.0 * u0 fifo.clear() for _ in range(FIFO_LEN): ff.run(si, bias) fifo.append(ff.obs.copy()[0:12]) save_states_arr = np.array(fifo, dtype=DATA_TYPE) # Save DDF+FIFO checkpoint for replay (state right after warmup, before step 0) ff.get_ddf() np.save(os.path.join(out_dir, "ddf_checkpoint.npy"), ff.ddf) np.save(os.path.join(out_dir, "fifo_checkpoint.npy"), save_states_arr) ff.apply_ddf() # === PPO inference === print("=== PPO inference ===") model = load_ppo_model(model_path_for_scene(scene_name), device=f"cuda:{device_id}", s_dim=s_dim, a_dim=3) model.set_random_seed(19) fifo = deque(maxlen=FIFO_LEN) for s in save_states_arr: fifo.append(np.array(s, dtype=DATA_TYPE)) obs = np.zeros(s_dim, dtype=np.float32) sens_c, forc_c, act_c, rew_c, sim_c = [], [], [], [], [] for step in range(n_steps): action, _ = model.predict(obs, deterministic=True) action = action.astype(np.float32).flatten() act_c.append(action.copy()) temp_a = np.zeros(n_env, dtype=DATA_TYPE) omega = (action * ac_scale + np.array(ac_bias, dtype=np.float32)) * u0 temp_a[3:6] = omega ff.context.push() ff.run(si, temp_a) ff.context.pop() obs_slice = ff.obs.copy()[0:12] fifo.append(obs_slice) sens_c.append(obs_slice[0:6]) forc_c.append(obs_slice[6:12]) # obs dimension depends on model type: # d1a3o12_* = 12-dim (forces + sens only) # d1a3o14_* = 14-dim (forces + sens + target_cd + target_cl) forces_norm = obs_slice[6:12] / force_norm_fact sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact if s_dim == 14: target_recon = gen_target_states_at(step, target_harmonics) t_cd_n = float(target_recon[0]) / force_norm_fact t_cl_n = float(target_recon[1]) / force_norm_fact obs = np.clip(np.hstack([forces_norm, sens_norm, t_cd_n, t_cl_n]), -1.0, 1.0).astype(np.float32) else: obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32) # Reward sarr = np.array(fifo, dtype=np.float32) if len(sarr) >= CONV_LEN: f = sarr[-1, 6:12] / force_norm_fact cd = float(f[0] + f[2] + f[4]) cl = float(f[1] + f[3] + f[5]) # DTW ref_seq = target_states[CONV_LEN:2*CONV_LEN, 3] cur_seq = sarr[-CONV_LEN:, 1] lag = calc_lag(ref_seq, cur_seq) sim_sum = 0.0 for i in range(6): t_seq = np.roll(target_states[:, i+2], -lag)[CONV_LEN:2*CONV_LEN] s_seq = sarr[-CONV_LEN:, i] sim_sum += calc_dtw_sim(t_seq, s_seq) / 6.0 similarities = float(sim_sum) sim_c.append(similarities) t_recon = gen_target_states_at(step, target_harmonics) t_cd = float(t_recon[0]) / force_norm_fact t_cl = float(t_recon[1]) / force_norm_fact r_cd = np.exp(-abs((cd - t_cd) * 10)) r_cl = np.exp(-abs((cl - t_cl) * 10)) r_sim = np.exp(-10 * abs(similarities - 1)) reward = float(min(0.3*r_cd + 0.3*r_cl + 0.4*r_sim, 1.0)) rew_c.append(reward) sens_arr = np.array(sens_c, dtype=np.float32) forc_arr = np.array(forc_c, dtype=np.float32) act_arr = np.array(act_c, dtype=np.float32) np.savez(os.path.join(out_dir, "controlled.npz"), sensors=sens_arr, forces=forc_arr, actions=act_arr, rewards=np.array(rew_c, dtype=np.float32)) save_vorticity_png(os.path.join(out_dir, "vorticity_controlled.png"), vorticity_from_ddf(ff, u0=u0), title=f"{scene_name} controlled") tail = min(100, len(rew_c)) avg_reward = float(np.mean(rew_c[-tail:])) if tail > 0 else 0.0 avg_sim = float(np.mean(sim_c[-tail:])) if sim_c else 0.0 print(f" reward={avg_reward:.4f} similarity={avg_sim:.4f}") result = {"scene": scene_name, "similarity": avg_sim, "avg_reward": avg_reward} with open(os.path.join(out_dir, "result.json"), "w") as f: json.dump(result, f, indent=2) del ff, model return result def main(): ap = argparse.ArgumentParser() ap.add_argument("--scene", type=str, default="illusion_1.0L", help="Scene name (illusion_0.75L, illusion_1.0L, illusion_1.5L)") ap.add_argument("--diameter", type=float, default=None, help="Diameter shortcut (0.75, 1.0, 1.5)") ap.add_argument("--device", type=int, default=2) ap.add_argument("--steps", type=int, default=200) args = ap.parse_args() if args.diameter is not None: scene_name = f"illusion_{args.diameter}L" else: scene_name = args.scene if scene_name not in get_scene_list("illusion"): print(f"Unknown scene: {scene_name}. Available: {get_scene_list('illusion')}") return 1 t0 = time.time() r = run_single(scene_name, args.device, args.steps) print(f"Done in {time.time()-t0:.1f}s: sim={r['similarity']:.4f}") if __name__ == "__main__": sys.exit(main())