351 lines
13 KiB
Python
351 lines
13 KiB
Python
# analysis_crossre/scripts/diagnose_equivariance.py
|
|
"""Phase A2-A3: diagnose PPO control-law equivariance under G operator.
|
|
|
|
Usage::
|
|
|
|
conda run -n pycuda_3_10 python diagnose_equivariance.py --re 100 --device 0
|
|
|
|
conda run -n pycuda_3_10 python diagnose_equivariance.py --re all --device 0
|
|
|
|
Output per Re: ``output/analysis_crossre/diagnostic/equivariance_re{re}.json``
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
_PROJ = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
|
if _PROJ not in sys.path:
|
|
sys.path.insert(0, _PROJ)
|
|
from LegacyCelerisLab import FlowField # noqa: E402
|
|
from LegacyCelerisLab import utils as legacy_utils # noqa: E402
|
|
|
|
from utils import (
|
|
action_to_physical,
|
|
compute_dimensionless,
|
|
apply_G_x,
|
|
apply_G_alpha,
|
|
load_ppo_model,
|
|
nu_from_re,
|
|
load_legacy_configs,
|
|
build_karman_cloak_env,
|
|
add_pinball,
|
|
build_observation,
|
|
scale_action,
|
|
)
|
|
from cfg import (
|
|
CONFIG_DIR,
|
|
OUTPUT_DIR,
|
|
MODEL_DIR,
|
|
SAMPLE_INTERVAL,
|
|
FIFO_LEN,
|
|
CONV_LEN,
|
|
S_DIM,
|
|
A_DIM,
|
|
ACTION_SCALE,
|
|
ACTION_BIAS,
|
|
U0,
|
|
RE_CASES_TRAIN,
|
|
RE_LABEL_MAP,
|
|
)
|
|
|
|
DATA_TYPE = np.float32
|
|
|
|
|
|
def diagnose_one_re(re_code: int, ppo_device: int, cfd_device: int, output_root: str) -> dict:
|
|
"""Run equivariance diagnosis for one Re case."""
|
|
os.makedirs(output_root, exist_ok=True)
|
|
|
|
nu = nu_from_re(re_code, u0=U0)
|
|
mu = 2.0 / re_code
|
|
label = RE_LABEL_MAP.get(re_code, f"Re{re_code}")
|
|
print(f"\n{'='*60}")
|
|
print(f"Diagnosing: {label} nu={nu:.6f} mu={mu:.6f}")
|
|
print(f"{'='*60}")
|
|
|
|
# Build full environment (dist + sensors + pinball)
|
|
cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
|
|
field_cfg = field_cfg._replace(viscosity=float(nu))
|
|
ff = FlowField(field_cfg, cuda_cfg, device_id=cfd_device)
|
|
|
|
# Stabilize and get to controlled state
|
|
target_states, _ = build_karman_cloak_env(
|
|
ff, u0=U0, l0=20.0, sample_interval=SAMPLE_INTERVAL,
|
|
fifo_len=FIFO_LEN, data_type=DATA_TYPE,
|
|
)
|
|
norm = add_pinball(
|
|
ff, l0=20.0, u0=U0, sample_interval=SAMPLE_INTERVAL,
|
|
fifo_len=FIFO_LEN, data_type=DATA_TYPE,
|
|
action_bias=ACTION_BIAS,
|
|
)
|
|
|
|
# Load PPO model
|
|
model_path = None
|
|
for rc, mn in RE_CASES_TRAIN:
|
|
if rc == re_code:
|
|
model_path = os.path.join(MODEL_DIR, "old", f"{mn}.zip")
|
|
break
|
|
if model_path is None or not os.path.isfile(model_path):
|
|
return {"re_code": re_code, "error": f"No model for Re{re_code}"}
|
|
|
|
model = load_ppo_model(model_path, device=f"cuda:{ppo_device}")
|
|
model.set_random_seed(0)
|
|
|
|
# Collect rollout data with PPO
|
|
ff.restore_ddf()
|
|
ff.apply_ddf()
|
|
|
|
# Bias FIFO
|
|
bias_action = scale_action(
|
|
np.zeros(3, dtype=np.float32),
|
|
scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0, n_total_bodies=7,
|
|
)
|
|
from collections import deque
|
|
fifo = deque(maxlen=FIFO_LEN)
|
|
for _ in range(FIFO_LEN):
|
|
ff.context.push()
|
|
try:
|
|
ff.run(SAMPLE_INTERVAL, bias_action)
|
|
finally:
|
|
ff.context.pop()
|
|
fifo.append(ff.obs.copy()[2:14])
|
|
|
|
n_steps = 150
|
|
obs_hist = np.zeros((n_steps, 12), dtype=np.float64)
|
|
alpha_hist = np.zeros((n_steps, 3), dtype=np.float64)
|
|
obs = np.zeros(S_DIM, dtype=np.float32)
|
|
|
|
for step in range(n_steps):
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
action = action.astype(np.float32).flatten()
|
|
|
|
# Convert to physical
|
|
action_arr = scale_action(
|
|
action, scale=ACTION_SCALE, bias=ACTION_BIAS,
|
|
u0=U0, n_total_bodies=7,
|
|
)
|
|
ff.context.push()
|
|
try:
|
|
ff.run(SAMPLE_INTERVAL, action_arr)
|
|
finally:
|
|
ff.context.pop()
|
|
|
|
obs_slice = ff.obs.copy()[2:14]
|
|
fifo.append(obs_slice)
|
|
alpha = action_to_physical(
|
|
action.reshape(1, 3), scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
|
|
).flatten()
|
|
|
|
obs_hist[step] = obs_slice
|
|
alpha_hist[step] = alpha
|
|
obs = build_observation(obs_slice, norm)
|
|
|
|
del ff
|
|
|
|
# ---- Equivariance diagnosis ----
|
|
dim = compute_dimensionless(obs_hist[:, 0:6], obs_hist[:, 6:12], u0=U0, d=20.0)
|
|
|
|
# Compute memory terms
|
|
a_prev = np.zeros_like(alpha_hist)
|
|
a_prev2 = np.zeros_like(alpha_hist)
|
|
a_prev[1:] = alpha_hist[:-1]
|
|
a_prev2[2:] = alpha_hist[:-2]
|
|
|
|
# Diagnostic 1: front bias check (mean of alpha_F)
|
|
mean_alpha_F = float(np.mean(alpha_hist[:, 0]))
|
|
std_alpha_F = float(np.std(alpha_hist[:, 0]))
|
|
front_bias_score = abs(mean_alpha_F) / (std_alpha_F + 1e-12)
|
|
|
|
# Diagnostic 2: check front equivariance
|
|
# For each point, compute PPO(Gx) by feeding G-transformed obs through model
|
|
eq_front_errors = []
|
|
eq_exchange_b_errors = []
|
|
eq_exchange_t_errors = []
|
|
eq_front_noise_floor = []
|
|
|
|
for t in range(2, n_steps):
|
|
# Get original obs and Gx
|
|
Gx = apply_G_x(
|
|
dim["u_hat_B"][t:t+1], dim["u_hat_C"][t:t+1], dim["u_hat_T"][t:t+1],
|
|
dim["v_hat_B"][t:t+1], dim["v_hat_C"][t:t+1], dim["v_hat_T"][t:t+1],
|
|
dim["Cd_F"][t:t+1], dim["Cd_T"][t:t+1], dim["Cd_B"][t:t+1],
|
|
dim["Cl_F"][t:t+1], dim["Cl_T"][t:t+1], dim["Cl_B"][t:t+1],
|
|
a_prev[t:t+1, 0], a_prev[t:t+1, 2], a_prev[t:t+1, 1],
|
|
a_prev2[t:t+1, 0] - a_prev[t:t+1, 0],
|
|
a_prev2[t:t+1, 2] - a_prev[t:t+1, 2],
|
|
a_prev2[t:t+1, 1] - a_prev[t:t+1, 1],
|
|
)
|
|
|
|
# Build Gx observation for PPO: we need the normalized obs
|
|
# The Gx in raw sensor/force space requires inverting the dimensionless transform
|
|
# Actually easier: compute what PPO would predict for the G state
|
|
# by transforming the raw obs and feeding it
|
|
|
|
# Build raw obs corresponding to Gx
|
|
raw_Gx = np.zeros(12, dtype=np.float64)
|
|
# Sensors: reorder + sign flip
|
|
# Original raw: [s0_ux, s0_uy, s1_ux, s1_uy, s2_ux, s2_uy] = top, center, bottom
|
|
# G: bottom->top, center->center, top->bottom
|
|
raw_Gx[0] = obs_hist[t, 4] # s0_ux <- s2_ux (bottom -> top, streamwise no sign)
|
|
raw_Gx[1] = -obs_hist[t, 5] # s0_uy <- -s2_uy (bottom -> top, cross sign flip)
|
|
raw_Gx[2] = obs_hist[t, 2] # s1_ux maintains (center)
|
|
raw_Gx[3] = -obs_hist[t, 3] # s1_uy = -s1_uy (center cross sign flip)
|
|
raw_Gx[4] = obs_hist[t, 0] # s2_ux <- s0_ux (top -> bottom)
|
|
raw_Gx[5] = -obs_hist[t, 1] # s2_uy <- -s0_uy (top -> bottom, cross sign flip)
|
|
# Forces: reorder + sign
|
|
# ordering: [front_fx, front_fy, bottom_fx, bottom_fy, top_fx, top_fy]
|
|
# G: front_fx -> front_fx (no sign), front_fy -> -front_fy
|
|
# bottom <-> top
|
|
raw_Gx[6] = obs_hist[t, 6] # front_fx unchanged
|
|
raw_Gx[7] = -obs_hist[t, 7] # front_fy sign flip
|
|
raw_Gx[8] = obs_hist[t, 10] # bottom_fx <- top_fx
|
|
raw_Gx[9] = -obs_hist[t, 11] # bottom_fy <- -top_fy
|
|
raw_Gx[10] = obs_hist[t, 8] # top_fx <- bottom_fx
|
|
raw_Gx[11] = -obs_hist[t, 9] # top_fy <- -bottom_fy
|
|
|
|
# Build normalized PPO observation from Gx
|
|
obs_Gx = build_observation(raw_Gx, norm)
|
|
|
|
# Predict action for Gx
|
|
action_Gx, _ = model.predict(obs_Gx, deterministic=True)
|
|
action_Gx = action_Gx.astype(np.float32).flatten()
|
|
alpha_Gx = action_to_physical(
|
|
action_Gx.reshape(1, 3), scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
|
|
).flatten()
|
|
|
|
# What equivariance says Gx should produce (with CORRECTED G)
|
|
# G([aF, aT, aB]) = [-aF, -aB, -aT]
|
|
alpha_Gx_expected = apply_G_alpha(alpha_hist[t])
|
|
|
|
# Front error: PPO(Gx)[0] should == G(PPO(x))[0] = -aF(x)
|
|
eq_front_errors.append(abs(float(alpha_Gx[0]) - float(alpha_Gx_expected[0])))
|
|
|
|
# Rear error (CORRECTED): PPO(Gx)[1] should == G(PPO(x))[1] = -aT(x)
|
|
# PPO(Gx)[2] should == G(PPO(x))[2] = -aB(x)
|
|
# Previously this incorrectly checked alpha_B(x) == alpha_T(Gx)
|
|
eq_exchange_b_errors.append(abs(float(alpha_Gx[1]) - float(alpha_Gx_expected[1])))
|
|
eq_exchange_t_errors.append(abs(float(alpha_Gx[2]) - float(alpha_Gx_expected[2])))
|
|
|
|
# Noise floor: difference between same-state replicate predictions
|
|
# (we approximate by checking prediction consistency)
|
|
action2, _ = model.predict(obs_Gx, deterministic=True)
|
|
alpha_Gx2 = action_to_physical(
|
|
action2.reshape(1, 3), scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
|
|
).flatten()
|
|
eq_front_noise_floor.append(abs(float(alpha_Gx2[0]) - float(alpha_Gx[0])))
|
|
|
|
eq_front_errors = np.array(eq_front_errors)
|
|
eq_exchange_b = np.array(eq_exchange_b_errors)
|
|
eq_exchange_t = np.array(eq_exchange_t_errors)
|
|
eq_noise = np.array(eq_front_noise_floor)
|
|
|
|
# Scale equivariance errors by action range for relative measure
|
|
alpha_range = float(np.max(np.abs(alpha_hist[2:])))
|
|
rel_front_err = float(np.mean(eq_front_errors) / (alpha_range + 1e-12))
|
|
rel_exchange_b_err = float(np.mean(eq_exchange_b) / (alpha_range + 1e-12))
|
|
rel_exchange_t_err = float(np.mean(eq_exchange_t) / (alpha_range + 1e-12))
|
|
|
|
# Combined rear error (max of bottom and top)
|
|
rel_exchange_err = max(rel_exchange_b_err, rel_exchange_t_err)
|
|
|
|
# Diagnostic 3: cross-correlation between alpha_T and -alpha_B
|
|
if len(alpha_hist) > 10:
|
|
# After initial transient
|
|
tail = n_steps // 2
|
|
corr_TB = float(np.corrcoef(alpha_hist[tail:, 2], -alpha_hist[tail:, 1])[0, 1])
|
|
else:
|
|
corr_TB = float("nan")
|
|
|
|
result = {
|
|
"re_code": re_code,
|
|
"mu": mu,
|
|
"n_samples": n_steps,
|
|
"alpha_range": alpha_range,
|
|
"front_bias": {
|
|
"mean_alpha_F": mean_alpha_F,
|
|
"std_alpha_F": std_alpha_F,
|
|
"bias_over_std": front_bias_score,
|
|
"bias_significant": front_bias_score > 2.0,
|
|
},
|
|
"equivariance_front": {
|
|
"mean_abs_error": float(np.mean(eq_front_errors)),
|
|
"max_abs_error": float(np.max(eq_front_errors)),
|
|
"relative_error": rel_front_err,
|
|
"noise_floor": float(np.mean(eq_noise)),
|
|
"signal_to_noise": float(np.mean(eq_front_errors) / (np.mean(eq_noise) + 1e-12)),
|
|
},
|
|
"equivariance_rear_bottom": {
|
|
"mean_abs_error": float(np.mean(eq_exchange_b)),
|
|
"max_abs_error": float(np.max(eq_exchange_b)),
|
|
"relative_error": rel_exchange_b_err,
|
|
},
|
|
"equivariance_rear_top": {
|
|
"mean_abs_error": float(np.mean(eq_exchange_t)),
|
|
"max_abs_error": float(np.max(eq_exchange_t)),
|
|
"relative_error": rel_exchange_t_err,
|
|
},
|
|
"top_bottom_correlation": {
|
|
"corr_alphaT_vs_negAlphaB": corr_TB,
|
|
},
|
|
"equivariance_verdict": "PASS" if (rel_front_err < 0.20 and rel_exchange_err < 0.20) else "REVIEW",
|
|
}
|
|
|
|
print(f" Front bias: mean_alpha_F={mean_alpha_F:.6f} |bias|/std={front_bias_score:.3f}")
|
|
print(f" Front equiv err: mean={np.mean(eq_front_errors):.6f} rel={rel_front_err:.3%}")
|
|
print(f" Rear-bot err: mean={np.mean(eq_exchange_b):.6f} rel={rel_exchange_b_err:.3%}")
|
|
print(f" Rear-top err: mean={np.mean(eq_exchange_t):.6f} rel={rel_exchange_t_err:.3%}")
|
|
print(f" T vs -B corr: {corr_TB:.4f}")
|
|
print(f" Verdict: {result['equivariance_verdict']}")
|
|
|
|
with open(os.path.join(output_root, f"equivariance_re{re_code}.json"), "w") as f:
|
|
json.dump(result, f, indent=2)
|
|
print(f" Saved to {output_root}/equivariance_re{re_code}.json")
|
|
|
|
return result
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Equivariance diagnosis for PPO cloak control")
|
|
ap.add_argument("--re", type=str, default="all",
|
|
help='Re case: 50,100,200,400, or "all"')
|
|
ap.add_argument("--device", type=int, default=0, help="GPU device for PPO model")
|
|
ap.add_argument("--cfd-device", type=int, default=2, help="GPU device for CFD simulation")
|
|
args = ap.parse_args()
|
|
|
|
if args.re.lower() == "all":
|
|
re_list = [rc for rc, _ in RE_CASES_TRAIN]
|
|
else:
|
|
re_list = [int(args.re)]
|
|
|
|
# Store device args for use in diagnose_one_re
|
|
device_id = args.device
|
|
cfd_device = args.cfd_device
|
|
|
|
diag_root = os.path.join(OUTPUT_DIR, "diagnostic")
|
|
os.makedirs(diag_root, exist_ok=True)
|
|
|
|
all_results = []
|
|
for re_code in re_list:
|
|
res = diagnose_one_re(re_code, device_id, cfd_device, diag_root)
|
|
all_results.append(res)
|
|
|
|
summary = {
|
|
"summary": {
|
|
"equivariance_verdicts": {r["re_code"]: r.get("equivariance_verdict", "ERROR")
|
|
for r in all_results}
|
|
},
|
|
"details": all_results,
|
|
}
|
|
with open(os.path.join(diag_root, "equivariance_summary.json"), "w") as f:
|
|
json.dump(summary, f, indent=2)
|
|
print(f"\nSummary saved to {diag_root}/equivariance_summary.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|