192 lines
6.6 KiB
Python
192 lines
6.6 KiB
Python
"""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,
|
|
}
|