第二轮：整理两个工作目录

2026-06-10 15:59:52 +08:00
parent d1b9922c6b
commit 096d9dcd0f
130 changed files with 13171 additions and 7263 deletions
@@ -0,0 +1,191 @@
+"""G-operator and equivariance tools.
+
+Provides G-operator transformations, dimensionless conversion,
+and equivariance diagnostics for PPO control laws.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+
+from .feature_builder import compute_dimensionless as _compute_dimless
+
+
+def apply_G_alpha(alpha: np.ndarray) -> np.ndarray:
+    """Apply mirror G to action: [aF, aT, aB] -> [-aF, -aB, -aT]."""
+    return np.array([-alpha[0], -alpha[2], -alpha[1]], dtype=alpha.dtype)
+
+
+def apply_G_raw(obs_slice: np.ndarray,
+                a_prev: np.ndarray,
+                a_prev2: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Apply G to raw obs slice [sensor(6)+force(6)] and action arrays.
+
+    Parameters
+    ----------
+    obs_slice : (12,) raw obs [s0_ux,uy, s1_ux,uy, s2_ux,uy, f0_fx,fy, f1_fx,fy, f2_fx,fy]
+    a_prev : (3,) physical omega at t-1
+    a_prev2 : (3,) physical omega at t-2
+
+    Returns
+    -------
+    G_obs : (12,) transformed obs slice
+    G_a_prev : (3,) transformed a_prev
+    G_a_prev2 : (3,) transformed a_prev2
+    """
+    G_obs = np.zeros(12, dtype=np.float64)
+    # sensors: swap top(0,1) <-> bottom(4,5), negate v
+    G_obs[0] = obs_slice[4]
+    G_obs[1] = -obs_slice[5]
+    G_obs[2] = obs_slice[2]
+    G_obs[3] = -obs_slice[3]
+    G_obs[4] = obs_slice[0]
+    G_obs[5] = -obs_slice[1]
+    # forces: swap bottom(2,3) <-> top(4,5), negate fy
+    G_obs[6] = obs_slice[6]
+    G_obs[7] = -obs_slice[7]
+    G_obs[8] = obs_slice[10]
+    G_obs[9] = -obs_slice[11]
+    G_obs[10] = obs_slice[8]
+    G_obs[11] = -obs_slice[9]
+
+    G_a_prev = np.array([-a_prev[0], -a_prev[2], -a_prev[1]], dtype=np.float64)
+    G_a_prev2 = np.array([-a_prev2[0], -a_prev2[2], -a_prev2[1]], dtype=np.float64)
+    return G_obs, G_a_prev, G_a_prev2
+
+
+def check_equivariance(
+    model: Any,
+    obs_slice_series: np.ndarray,  # (T, 12) raw obs
+    actions_phys: np.ndarray,      # (T, 3) physical omega
+    norm: dict,
+    action_scale: float = 8.0,
+    action_bias: Tuple[float, float, float] = (0.0, -4.0, 4.0),
+    u0: float = 0.01,
+) -> Dict[str, float]:
+    """Check G-equivariance of a PPO model over a time series.
+
+    Returns dict with front/rear equivariance errors.
+    """
+    from .cfd_interface import build_observation, action_to_physical
+
+    T = min(obs_slice_series.shape[0], actions_phys.shape[0])
+    ef, eb, et = [], [], []
+
+    for t in range(2, T):
+        # Get current obs
+        osl = obs_slice_series[t]
+        a_prev = actions_phys[t - 1] if t > 0 else actions_phys[0]
+        a_prev2 = actions_phys[t - 2] if t > 1 else actions_phys[0]
+
+        # Predict action for current state
+        obs = build_observation(osl, norm)
+        act, _ = model.predict(obs, deterministic=True)
+        act = act.astype(np.float32).flatten()
+        alpha = action_to_physical(act.reshape(1, 3),
+                                   scale=action_scale, bias=action_bias, u0=u0).flatten()
+
+        # Apply G to state
+        G_obs, _, _ = apply_G_raw(osl, a_prev, a_prev2)
+        obs_G = build_observation(G_obs, norm)
+        act_G, _ = model.predict(obs_G, deterministic=True)
+        act_G = act_G.astype(np.float32).flatten()
+        alpha_G = action_to_physical(act_G.reshape(1, 3),
+                                     scale=action_scale, bias=action_bias, u0=u0).flatten()
+
+        # Expected: G(alpha) = [-aF, -aB, -aT]
+        expected = apply_G_alpha(alpha)
+
+        ef.append(abs(float(alpha_G[0]) - float(expected[0])))
+        eb.append(abs(float(alpha_G[1]) - float(expected[1])))
+        et.append(abs(float(alpha_G[2]) - float(expected[2])))
+
+    ef_arr = np.array(ef)
+    eb_arr = np.array(eb)
+    et_arr = np.array(et)
+    alpha_range = float(np.max(np.abs(actions_phys[2:])))
+
+    return {
+        "front_mean_abs_error": float(np.mean(ef_arr)),
+        "front_rel_error": float(np.mean(ef_arr) / (alpha_range + 1e-12)),
+        "rear_bottom_rel_error": float(np.mean(eb_arr) / (alpha_range + 1e-12)),
+        "rear_top_rel_error": float(np.mean(et_arr) / (alpha_range + 1e-12)),
+        "alpha_range": alpha_range,
+    }
+
+
+def diagnose_one_re(model, ff, target_states, norm, config, n_steps=150) -> dict:
+    """Run PPO inference and check equivariance.
+
+    Parameters
+    ----------
+    model : loaded PPO model
+    ff : FlowField instance (must be at saved checkpoint state)
+    target_states : (FIFO_LEN, 6) target sensor signals
+    norm : norm dict
+    config : scene config dict with action_scale, action_bias, u0, etc.
+
+    Returns
+    -------
+    dict with equivariance metrics.
+    """
+    from collections import deque
+    from .cfd_interface import (build_observation, scale_action,
+                                 action_to_physical, compute_similarity)
+
+    action_scale = config.get("action_scale", 8.0)
+    action_bias = config.get("action_bias", (0.0, -4.0, 4.0))
+    u0 = config.get("u0", 0.01)
+    sample_interval = config.get("sample_interval", 800)
+    fifo_len = config.get("fifo_len", 150)
+    n_obj_total = config.get("n_objects_total", 7)
+
+    ff.restore_ddf()
+    ff.apply_ddf()
+
+    # Bias FIFO init
+    fifo = deque(maxlen=fifo_len)
+    bias_arr = scale_action(np.zeros(3, dtype=np.float32),
+                            scale=action_scale, bias=action_bias,
+                            u0=u0, n_total_bodies=n_obj_total)
+    for _ in range(fifo_len):
+        ff.run(sample_interval, bias_arr)
+        fifo.append(ff.obs.copy()[2:14])
+
+    # Inference
+    obs_array = []
+    action_array = []
+    obs = np.zeros(12, dtype=np.float32)
+
+    for _ in range(n_steps):
+        act, _ = model.predict(obs, deterministic=True)
+        act = act.astype(np.float32).flatten()
+        action_array.append(act.copy())
+
+        action_arr = scale_action(act, scale=action_scale, bias=action_bias,
+                                  u0=u0, n_total_bodies=n_obj_total)
+        ff.context.push()
+        ff.run(sample_interval, action_arr)
+        ff.context.pop()
+
+        obs_slice = ff.obs.copy()[2:14]
+        fifo.append(obs_slice)
+        obs_array.append(obs_slice)
+        obs = build_observation(obs_slice, norm)
+
+    obs_series = np.array(obs_array, dtype=np.float64)
+    actions_phys = action_to_physical(np.array(action_array),
+                                      scale=action_scale, bias=action_bias, u0=u0)
+    states_arr = np.array(list(fifo), dtype=np.float32)
+    sim = compute_similarity(target_states, states_arr[:, 0:6],
+                             config.get("conv_len", 30))
+
+    # Equivariance check
+    eq = check_equivariance(model, obs_series, actions_phys, norm,
+                            action_scale, action_bias, u0)
+
+    return {
+        "similarity": sim,
+        "equivariance": eq,
+    }