"""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, }