第二轮：整理两个工作目录
This commit is contained in:
Frank14f
@@ -0,0 +1,233 @@
|
||||
"""1L Illusion DRL inference (2U=0.02).
|
||||
|
||||
Usage:
|
||||
conda run -n pycuda_3_10 python scripts/collect_illusion.py --device 2 --steps 200
|
||||
|
||||
Output: data/illusion/illusion_1L/
|
||||
"""
|
||||
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)
|
||||
_ANALYSIS = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
if _ANALYSIS not in sys.path:
|
||||
sys.path.insert(0, _ANALYSIS)
|
||||
|
||||
from LegacyCelerisLab import FlowField
|
||||
|
||||
from CCD_analysis.configs import get_scene, 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)
|
||||
ff_tgt.add_cylinder((20.0 * L0, CENTER_Y, 0.0), 1.0 * L0)
|
||||
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}")
|
||||
|
||||
# Bias FIFO (matches legacy_env_imit: [0,0,0,0,-1*U0,1*U0])
|
||||
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)
|
||||
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])
|
||||
|
||||
# 14-dim obs
|
||||
forces_norm = obs_slice[6:12] / force_norm_fact
|
||||
sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact
|
||||
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)
|
||||
|
||||
# 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("--device", type=int, default=2)
|
||||
ap.add_argument("--steps", type=int, default=200)
|
||||
args = ap.parse_args()
|
||||
|
||||
t0 = time.time()
|
||||
r = run_single("illusion_1L", args.device, args.steps)
|
||||
print(f"Done in {time.time()-t0:.1f}s: sim={r['similarity']:.4f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user