diff --git a/.gitignore b/.gitignore index 594ec0e..481c963 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,10 @@ saved_models/ # Ignore visualization outputs figures/ video_frames/ + +# Output of koopmanrl_utils/skvi_policy_checks.py, skvi_sensitivity_checks.py and +# koopman_prediction_validation.py +skvi_policy_checks_results/ +skvi_sensitivity_checks_results/ +koopman_prediction_validation_results/ +koopmanrl_utils/__pycache__/ diff --git a/koopmanrl_utils/AGENTS.md b/koopmanrl_utils/AGENTS.md index 7b5be4c..9ce4392 100644 --- a/koopmanrl_utils/AGENTS.md +++ b/koopmanrl_utils/AGENTS.md @@ -29,6 +29,7 @@ koopmanrl_utils/ ├── dataframe_creator.py # Converts Tensorboard results to JSON data frames ├── interpret_koopman.json # Configuration of test file to test interpretability on ├── interpret_koopman.py # Ingests a Koopman configuration, and interprets its tensor +├── koopman_prediction_validation.py # Held-out one-step and multi-step prediction error of the Koopman tensor ├── plot_csv_from_tensorboards.py # Ingests Tensorboard results and generates csv files ├── process_episodic_returns.py # Generates episodic return plots from JSON dataframe ├── process_sakc_ablations.py # Generates the ablation plots for the Soft Actor Koopman-Critic from the JSON dataframes @@ -36,6 +37,8 @@ koopmanrl_utils/ ├── run_optimized_experiments.py # Runs the optimized experimental configurations of the Koopman algorithms ├── run_sakc_optimization.py # Runs the Soft Actor Koopman Critic hyperparameter optimization ├── run_skvi_optimization.py # Runs the Soft Koopman Value Iteration hyperparameter optimization +├── skvi_policy_checks.py # Reads the SKVI policy off the Koopman tensor and checks it (LQR, pruning) +├── skvi_sensitivity_checks.py # Koopman-tensor accuracy along the SKVI policy; sensitivity of SKVI's control └── tsne_koopman_tensor.py # Loads the saved Koopman tensors and generates a t-SNE plot from it ``` diff --git a/koopmanrl_utils/koopman_prediction_validation.py b/koopmanrl_utils/koopman_prediction_validation.py new file mode 100644 index 0000000..a9fa727 --- /dev/null +++ b/koopmanrl_utils/koopman_prediction_validation.py @@ -0,0 +1,434 @@ +"""Prediction error of the fitted Koopman tensor on held-out random-agent data. + +For each benchmark the Koopman tensor is identified from random-agent trajectories with the tuned SKVI dictionary +orders and identification budget of `configurations/`, on 80% of the trajectories, and evaluated on the remaining 20%. +The split is by whole trajectory, so no held-out transition comes from a trajectory used for the fit. Everything is +repeated over seeds and reported as the median (and quartiles) across seeds. + +One-step error (`onestep_error.csv`, `onestep_by_degree*.csv`, `onestep_table.tex`) + SKVI and SAKC use the tensor only through the one-step prediction K^u phi(x) of the dictionary; neither iterates + it nor decodes it back to the state. The error is therefore measured in dictionary space, as the root mean square + over transitions n and nonconstant dictionary functions j of + + (K^{u_n} phi(x_n) - E[phi(x'_n) | x_n, u_n])_j / s_j, + + where s_j is the root mean square of the target for function j, so that high-order monomials do not dominate. For + the deterministic benchmarks the conditional mean is phi(x'_n). For the double well it is a Monte-Carlo average + over the one-step noise, and the error of a single realisation phi(x'_n) about it is reported as the noise floor. + Because the time step is small, x' is close to x, so the same error with phi(x_n) in place of the prediction (the + persistence baseline) and the ratio of the two are reported too: the ratio is the fraction of the one-step change + of the dictionary that the model misses. The by-degree files give the held-out error per monomial degree. The + error of the decoded state, normalised by the root mean square of the target state, is kept in the unprefixed + columns for reference. + +Multi-step error (`multistep_.csv`) + The error of the predicted state along fresh trajectories under the recorded actions, against the horizon, for + two ways of rolling the model forward: iterating phi_{k+1} = K^{u_k} phi_k and decoding with B ("lifted"), and + re-evaluating the dictionary at the predicted state after every step, x_{k+1} = B^T K^{u_k} phi(x_k) + ("reproj"). For the double well the reference is a Monte-Carlo mean trajectory. This is a diagnostic of the model: + the algorithms never roll it forward. + +Training against held-out error (`train_vs_heldout_.csv`) + The one-step error on the training and on the held-out transitions against the number of identification + transitions, on nested subsamples (each smaller set is contained in the next). With the number of parameters + fixed, the least-squares training error rises with the sample size towards the approximation floor while the + held-out error falls towards it; overfitting would show as a gap that does not close. + +Usage (from the repository root): + + uv run -m koopmanrl_utils.koopman_prediction_validation # four benchmarks, seeds 123-130 + uv run -m koopmanrl_utils.koopman_prediction_validation --seeds 2 --seed0 0 # two seeds starting at 0 + +The files are written into `--output_dir` (default `koopman_prediction_validation_results/`, not tracked by git). +Random draws use NumPy's global generator and a per-seed `RandomState`, seeded in a fixed order, so a run is +reproducible from its seed up to floating-point rounding (differences of order 1e-12 between runs, which reach the +leading digits only for the linear system, whose errors are at machine precision). + +Reference: the results are in "Koopman-Assisted Reinforcement Learning" (Rozwood, Mehrez, Paehler, Sun and Brunton), +section "Evaluation" (one-step error), and in its electronic supplementary material, section "Koopman model +validation: additional results" (multi-step error, training against held-out error). +""" + +import argparse +import os + +import gym +import numpy as np +import torch + +torch.set_default_dtype(torch.float64) + +import koopmanrl.environments # noqa: F401,E402 (import registers the gym envs) +from koopmanrl.koopman_tensor.observables import torch_observables as obs # noqa: E402 +from koopmanrl.koopman_tensor.observables.torch_observables import ( # noqa: E402 + allMonomialPowers, +) +from koopmanrl.koopman_tensor.torch_tensor import KoopmanTensor, Regressor # noqa: E402 + +# Tuned SKVI tensor settings, as in configurations/skvi__hparams.json (num-paths, num-steps-per-path, +# state-order, action-order). +CONFIGS = { + "LinearSystem-v0": dict(label="Linear", num_paths=75, steps=250, so=2, ao=3, stochastic=False), + "Lorenz-v0": dict(label="Lorenz", num_paths=150, steps=250, so=3, ao=1, stochastic=False), + "FluidFlow-v0": dict(label="FluidFlow", num_paths=200, steps=225, so=4, ao=2, stochastic=False), + "DoubleWell-v0": dict(label="DoubleWell", num_paths=175, steps=100, so=2, ao=4, stochastic=True), +} + +HORIZON = 500 # rollout horizon (steps) +N_ROLLOUT_PATHS = 40 # fresh held-out trajectories used for the rollout test +N_MC = 512 # Monte-Carlo samples for the stochastic mean-trajectory reference +N_ID_GRID = [250, 500, 1000, 2000, 4000, 8000, 12000] +HELDOUT_FRAC = 0.2 +CLIP = 1e6 # guard against polynomial-model blow-up in long rollouts + + +# --------------------------------------------------------------------------- # +# Data +# --------------------------------------------------------------------------- # +def make_env(env_id, seed): + env = gym.make(env_id).unwrapped # unwrap TimeLimit; we control the loop length + env.action_space.seed(seed) + return env + + +def collect_paths(env, num_paths, steps): + """Random-agent rollouts; returns numpy arrays shaped (P, steps, dim).""" + sd = env.observation_space.shape[0] + ad = env.action_space.shape[0] + X = np.zeros((num_paths, steps, sd)) + Y = np.zeros((num_paths, steps, sd)) + U = np.zeros((num_paths, steps, ad)) + for p in range(num_paths): + s = np.asarray(env.reset(), dtype=float) + for t in range(steps): + a = np.asarray(env.action_space.sample(), dtype=float) + X[p, t], U[p, t] = s, a + s = np.asarray(env.step(a)[0], dtype=float) + Y[p, t] = s + return X, Y, U + + +def flat(A): + """(P, steps, dim) -> (dim, P*steps).""" + P, T, d = A.shape + return A.reshape(P * T, d).T + + +# --------------------------------------------------------------------------- # +# DoubleWell SDE helpers (vectorised; checked against the env at start-up) +# --------------------------------------------------------------------------- # +def dw_drift(Xd, Ud): + """Vectorised DoubleWell drift; Xd (2, N), Ud (1, N).""" + return np.stack([4 * Xd[0] - 4 * Xd[0] ** 3 + Ud[0], -2 * Xd[1] + Ud[0]]) + + +def dw_selfcheck(env): + rng = np.random.RandomState(0) + for _ in range(20): + x, u = rng.uniform(-2, 2, 2), rng.uniform(-1, 1, 1) + ref = np.asarray(env.continuous_f(u)(0.0, x), dtype=float).ravel() + assert np.allclose(ref, dw_drift(x[:, None], u[:, None])[:, 0]), "dw_drift out of sync with env" + + +def dw_mc_mean_trajectory(x0, Useq, dt, rng): + """Monte-Carlo E[x_h | x_0, u_0..u_{h-1}] by Euler-Maruyama (same scheme as the env). + x0 (2, P), Useq (P, H, 1) -> (2, P, H).""" + P, H, _ = Useq.shape + Xc = np.repeat(x0[:, :, None], N_MC, axis=2) # (2, P, M) + out = np.empty((2, P, H)) + for h in range(H): + u = Useq[:, h, 0][None, :, None] + drift = np.stack([4 * Xc[0] - 4 * Xc[0] ** 3 + u[0], -2 * Xc[1] + u[0]]) + w = rng.normal(size=(2, P, N_MC)) + # sigma(x) = [[0.7, x], [0, 0.5]] + diff = np.stack([0.7 * w[0] + Xc[0] * w[1], 0.5 * w[1]]) + Xc = np.clip(Xc + drift * dt + diff * np.sqrt(dt), -CLIP, CLIP) + out[:, :, h] = Xc.mean(axis=2) + return out + + +# --------------------------------------------------------------------------- # +# Model + metrics +# --------------------------------------------------------------------------- # +def rms(A): + return float(np.sqrt(np.mean(np.sum(A**2, axis=0)))) + + +def nrmse(pred, ref): + return rms(ref - pred) / rms(ref) + + +def build_tensor(Xd, Yd, Ud, so, ao): + return KoopmanTensor( + torch.tensor(Xd), + torch.tensor(Yd), + torch.tensor(Ud), + phi=obs.monomials(so), + psi=obs.monomials(ao), + regressor=Regressor("ols"), + ) + + +def predict(tensor, Xd, Ud): + out = np.asarray(tensor.f(torch.tensor(np.ascontiguousarray(Xd)), torch.tensor(np.ascontiguousarray(Ud)))) + return np.clip(np.nan_to_num(out, nan=CLIP, posinf=CLIP, neginf=-CLIP), -CLIP, CLIP) + + +def phi_of(tensor, Xd): + return np.asarray(tensor.phi(torch.tensor(np.ascontiguousarray(Xd)))) + + +def predict_phi(tensor, Xd, Ud): + """One-step prediction in dictionary space, K^u phi(x): exactly the quantity SKVI and SAKC consume.""" + out = np.asarray(tensor.phi_f(torch.tensor(np.ascontiguousarray(Xd)), torch.tensor(np.ascontiguousarray(Ud)))) + return np.clip(np.nan_to_num(out, nan=CLIP, posinf=CLIP, neginf=-CLIP), -CLIP, CLIP) + + +def ref_phi(tensor, env, Xd, Ud, Yd, stoch, rng, n_mc=128): + """Reference E[phi(x') | x, u]. Deterministic: phi(x'). DoubleWell: Monte-Carlo over the + one-step Gaussian increment (note E[phi(x')] != phi(E[x']) for nonlinear phi).""" + if not stoch: + return phi_of(tensor, Yd) + mean = Xd + env.dt * dw_drift(Xd, Ud) + acc = 0.0 + for _ in range(n_mc): + w = rng.normal(size=Xd.shape) + diff = np.stack([0.7 * w[0] + Xd[0] * w[1], 0.5 * w[1]]) * np.sqrt(env.dt) + acc = acc + phi_of(tensor, mean + diff) + return acc / n_mc + + +def lifted_metrics(pred, ref, base, realised=None): + """Scaled dictionary-space errors. Each dictionary coordinate is divided by its RMS over the + reference set so that high-order monomials do not dominate; constant coordinates are dropped.""" + keep = ref.std(axis=1) > 1e-12 + sc = np.sqrt(np.mean(ref[keep] ** 2, axis=1, keepdims=True)) + e = lambda A: float(np.sqrt(np.mean((A[keep] / sc) ** 2))) # noqa: E731 + out = dict(err=e(ref - pred), persistence=e(ref - base), increment=e(ref - pred) / e(ref - base)) + out["noise_floor"] = e(ref - realised) if realised is not None else 0.0 + return out + + +def lifted_step(tensor, Phi, Ud): + """phi_{k+1} = K^{u_k} phi_k for a batch; Phi (d, P), Ud (adim, P).""" + K_u = np.asarray(tensor.K_(torch.tensor(np.ascontiguousarray(Ud)))) + if K_u.ndim == 2: + K_u = K_u[None] + out = np.einsum("pij,jp->ip", K_u, Phi) + return np.clip(np.nan_to_num(out, nan=CLIP, posinf=CLIP, neginf=-CLIP), -CLIP, CLIP) + + +# --------------------------------------------------------------------------- # +# One seed of one environment +# --------------------------------------------------------------------------- # +def run_env(env_id, cfg, seed): + np.random.seed(seed) + torch.manual_seed(seed) + rng = np.random.RandomState(seed) + env = make_env(env_id, seed) + stoch = cfg["stochastic"] + if stoch: + dw_selfcheck(env) + + X, Y, U = collect_paths(env, cfg["num_paths"], cfg["steps"]) + perm = rng.permutation(cfg["num_paths"]) + n_te = max(1, int(HELDOUT_FRAC * cfg["num_paths"])) + te, tr = perm[:n_te], perm[n_te:] + Xtr, Ytr, Utr = flat(X[tr]), flat(Y[tr]), flat(U[tr]) + Xte, Yte, Ute = flat(X[te]), flat(Y[te]), flat(U[te]) + + ref_tr = Xtr + env.dt * dw_drift(Xtr, Utr) if stoch else Ytr + ref_te = Xte + env.dt * dw_drift(Xte, Ute) if stoch else Yte + + tensor = build_tensor(Xtr, Ytr, Utr, cfg["so"], cfg["ao"]) + + # one-step error of the decoded state + pred_te = predict(tensor, Xte, Ute) + onestep = dict( + train=nrmse(predict(tensor, Xtr, Utr), ref_tr), + heldout=nrmse(pred_te, ref_te), + persistence=nrmse(Xte, ref_te), + increment=rms(ref_te - pred_te) / rms(ref_te - Xte), + noise_floor=nrmse(Yte, ref_te) if stoch else 0.0, + ) + + rphi_tr = ref_phi(tensor, env, Xtr, Utr, Ytr, stoch, rng) + rphi_te = ref_phi(tensor, env, Xte, Ute, Yte, stoch, rng) + L = lifted_metrics( + predict_phi(tensor, Xte, Ute), rphi_te, phi_of(tensor, Xte), phi_of(tensor, Yte) if stoch else None + ) + Ltr = lifted_metrics(predict_phi(tensor, Xtr, Utr), rphi_tr, phi_of(tensor, Xtr)) + onestep.update( + phi_train=Ltr["err"], + phi_heldout=L["err"], + phi_persistence=L["persistence"], + phi_increment=L["increment"], + phi_noise_floor=L["noise_floor"], + ) + + # held-out dictionary-space error by monomial degree + deg = np.asarray(allMonomialPowers(X.shape[2], cfg["so"])).sum(axis=0) + pphi_te, bphi_te = predict_phi(tensor, Xte, Ute), phi_of(tensor, Xte) + bydeg = np.full((4, 2), np.nan) + for d in range(1, cfg["so"] + 1): + k = deg == d + m = lifted_metrics(pphi_te[k], rphi_te[k], bphi_te[k]) + bydeg[d - 1] = [m["err"], m["increment"]] + onestep["bydeg"] = bydeg + + # multi-step error along fresh trajectories + Xr, Yr, Ur = collect_paths(env, N_ROLLOUT_PATHS, HORIZON) # fresh held-out trajectories + x0 = Xr[:, 0, :].T # (dim, P) + if stoch: + ref_traj = dw_mc_mean_trajectory(x0, Ur, env.dt, rng) # (dim, P, H) + else: + ref_traj = np.transpose(Yr, (2, 0, 1)) + Xcur = x0.copy() + Phi = np.asarray(tensor.phi(torch.tensor(np.ascontiguousarray(x0)))) + B = np.asarray(tensor.B) + multistep = np.empty((HORIZON, 3)) # reprojected, lifted, persistence + for h in range(HORIZON): + Uk = Ur[:, h, :].T + Xcur = predict(tensor, Xcur, Uk) + Phi = lifted_step(tensor, Phi, Uk) + ref_h = ref_traj[:, :, h] + multistep[h] = [nrmse(Xcur, ref_h), nrmse(np.clip(B.T @ Phi, -CLIP, CLIP), ref_h), nrmse(x0, ref_h)] + + # training against held-out error on nested subsamples + order = rng.permutation(Xtr.shape[1]) + sweep = np.full((len(N_ID_GRID), 2), np.nan) + for j, n_id in enumerate(N_ID_GRID): + if n_id > Xtr.shape[1]: + continue + idx = order[:n_id] + t = build_tensor(Xtr[:, idx], Ytr[:, idx], Utr[:, idx], cfg["so"], cfg["ao"]) + a = lifted_metrics(predict_phi(t, Xtr[:, idx], Utr[:, idx]), rphi_tr[:, idx], phi_of(t, Xtr[:, idx])) + b = lifted_metrics(predict_phi(t, Xte, Ute), rphi_te, phi_of(t, Xte)) + sweep[j] = [a["err"], b["err"]] + + return onestep, multistep, sweep, getattr(env, "dt", 1.0) + + +# --------------------------------------------------------------------------- # +# Output +# --------------------------------------------------------------------------- # +def q(a, axis=0): + """median, lower quartile, upper quartile across seeds.""" + return np.nanmedian(a, axis), np.nanpercentile(a, 25, axis), np.nanpercentile(a, 75, axis) + + +def write_csv(path, header, rows): + with open(path, "w") as f: + f.write(",".join(header) + "\n") + for r in rows: + f.write(",".join(f"{v:.6e}" if isinstance(v, float) else str(v) for v in r) + "\n") + print(f"wrote {path}") + + +def tex_num(v): + if v == 0: + return "---" + m, e = f"{v:.1e}".split("e") + return f"${m}\\times10^{{{int(e)}}}$" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--output_dir", default="koopman_prediction_validation_results") + ap.add_argument("--seeds", type=int, default=8, help="number of seeds") + ap.add_argument("--seed0", type=int, default=123, help="first seed; seeds are seed0, seed0 + 1, ...") + args = ap.parse_args() + os.makedirs(args.output_dir, exist_ok=True) + + one_rows, tex_rows, deg_rows = [], [], [] + for env_id, cfg in CONFIGS.items(): + label = cfg["label"] + print(f"\n=== {env_id} (orders s{cfg['so']}/a{cfg['ao']}, {cfg['num_paths']}x{cfg['steps']}) ===") + ones, multis, sweeps = [], [], [] + for k in range(args.seeds): + o, m, s, dt = run_env(env_id, cfg, args.seed0 + k) + ones.append(o), multis.append(m), sweeps.append(s) + print(f" seed {args.seed0 + k}: held-out {o['heldout']:.3e} increment-norm. {o['increment']:.3e}") + + keys = [ + "phi_train", + "phi_heldout", + "phi_persistence", + "phi_increment", + "phi_noise_floor", + "train", + "heldout", + "persistence", + "increment", + "noise_floor", + ] + med = {k_: float(np.median([o[k_] for o in ones])) for k_ in keys} + one_rows.append([label] + [med[k_] for k_ in keys]) + tex_rows.append(f"{label} & " + " & ".join(tex_num(med[k_]) for k_ in keys[:5]) + r" \\") + + D = np.nanmedian(np.stack([o["bydeg"] for o in ones]), axis=0) + deg_rows.append([label] + [float(D[d, 0]) for d in range(4)] + [float(D[d, 1]) for d in range(4)]) + + M = np.stack(multis) # (seeds, H, 3) + rows = [] + for h in range(HORIZON): + r = [h + 1, float((h + 1) * dt)] + for c in range(2): + a, lo, hi = q(M[:, h, c]) + r += [float(a), float(lo), float(hi)] + r.append(float(np.median(M[:, h, 2]))) + rows.append(r) + write_csv( + os.path.join(args.output_dir, f"multistep_{label}.csv"), + ["h", "t", "reproj", "reproj_lo", "reproj_hi", "lifted", "lifted_lo", "lifted_hi", "persist"], + rows, + ) + + S = np.stack(sweeps) # (seeds, grid, 2) + rows = [] + for j, n_id in enumerate(N_ID_GRID): + if np.all(np.isnan(S[:, j, 0])): + continue + (a, alo, ahi), (b, blo, bhi) = q(S[:, j, 0]), q(S[:, j, 1]) + rows.append([n_id, float(a), float(alo), float(ahi), float(b), float(blo), float(bhi)]) + write_csv( + os.path.join(args.output_dir, f"train_vs_heldout_{label}.csv"), + ["n_id", "train", "train_lo", "train_hi", "heldout", "heldout_lo", "heldout_hi"], + rows, + ) + + write_csv( + os.path.join(args.output_dir, "onestep_error.csv"), + [ + "env", + "phi_train", + "phi_heldout", + "phi_persistence", + "phi_increment", + "phi_noise_floor", + "train", + "heldout", + "persistence", + "increment", + "noise_floor", + ], + one_rows, + ) + write_csv( + os.path.join(args.output_dir, "onestep_by_degree.csv"), + ["env"] + [f"err_deg{d}" for d in range(1, 5)] + [f"inc_deg{d}" for d in range(1, 5)], + deg_rows, + ) + # same data, one row per degree (pgfplots-friendly) + write_csv( + os.path.join(args.output_dir, "onestep_by_degree_long.csv"), + ["degree"] + [r[0] for r in deg_rows], + [[d + 1] + [r[1 + d] for r in deg_rows] for d in range(4)], + ) + with open(os.path.join(args.output_dir, "onestep_table.tex"), "w") as f: + f.write("\n".join(tex_rows) + "\n\\bottomrule\n") # \bottomrule must live here: \input before it breaks tabular + print(f"wrote {os.path.join(args.output_dir, 'onestep_table.tex')}") + + +if __name__ == "__main__": + main() diff --git a/koopmanrl_utils/skvi_policy_checks.py b/koopmanrl_utils/skvi_policy_checks.py new file mode 100644 index 0000000..7f5b740 --- /dev/null +++ b/koopmanrl_utils/skvi_policy_checks.py @@ -0,0 +1,577 @@ +"""Read the deployed policy of Soft Koopman Value Iteration (SKVI) off the Koopman tensor, and check the reading. + +For each seed, SKVI is trained with the tuned configuration in `configurations/`. With the polynomial action +dictionary psi(u) = (1, u, u^2, ...), the fitted continuation w^T K^u phi(x) is a polynomial in the action whose +coefficients are functions of the state read off the slices of the tensor (`action_fields`), so the deployed Gibbs +policy can be written down in closed form (`policy_reading`). The checks compare that reading with references that +do not use the tensor. + +Linear system + The learned quadratic form of the value function is compared with the Riccati matrix of the discounted linear + quadratic regulator (LQR), and the gain of the policy's mean action with the discounted LQR gain. + +Double well + The value function is pruned in two ways (odd terms removed; diagonal quadratic only), and the returns of paired + rollouts and the total-variation distance between the pruned and the full policy are compared. The gain of the + mean action near the origin and the agreement of the closed-form Gaussian mean with the grid policy are reported. + +Usage (from the repository root): + + uv run -m koopmanrl_utils.skvi_policy_checks # both environments, seeds 100-107 + uv run -m koopmanrl_utils.skvi_policy_checks --environments LinearSystem --seeds 100 101 + uv run -m koopmanrl_utils.skvi_policy_checks --summarize_only --plot + +Each run writes one JSON file into `--output_dir` (default `skvi_policy_checks_results/`, not tracked by git). Random +draws use NumPy's global generator, seeded per run as in `koopmanrl.soft_koopman_value_iteration`, in a fixed order, so +a run is reproducible from its seed. + +Reference: the theory and the results of these checks are in the electronic supplementary material of +"Koopman-Assisted Reinforcement Learning" (Rozwood, Mehrez, Paehler, Sun and Brunton), section "Additional validation +and interpretability material". +""" + +import contextlib +import glob +import io +import json +import os +import time + +import gym +import numpy as np +import torch +from scipy.linalg import solve_discrete_are +from tap import Tap + +import koopmanrl.environments # noqa: F401 (registers the environments with gym) +from koopmanrl.koopman_observables import allMonomialPowers +from koopmanrl.koopman_observables import monomials as monomial_dictionary +from koopmanrl.soft_koopman_value_iteration import ( + DiscreteKoopmanValueIterationPolicy, + KoopmanTensor, + Regressor, + generate_koopman_tensor, +) + +# Settings shared with the SKVI script (koopmanrl/soft_koopman_value_iteration.py defaults). +GAMMA = 0.99 +ALPHA = 1.0 +NUM_ACTIONS = 101 +BATCH_SIZE = 2**14 +PROBABILITY_FLOOR = torch.finfo(torch.float64).eps # the `delta` added to every action probability by the package + +CONFIG_FILES = { + "LinearSystem-v0": "skvi_linear_system_hparams.json", + "DoubleWell-v0": "skvi_double_well_hparams.json", + "FluidFlow-v0": "skvi_fluid_flow_hparams.json", # used by skvi_sensitivity_checks + "Lorenz-v0": "skvi_lorenz_hparams.json", # used by skvi_sensitivity_checks +} +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +class ArgumentParser(Tap): + environments: list[str] = ["LinearSystem", "DoubleWell"] # environments to check + seeds: list[int] = list(range(100, 108)) # one SKVI run per seed + output_dir: str = "skvi_policy_checks_results" # one JSON file per run is written here + episodes: int = 10 # paired rollout episodes per variant (double well) + horizon: int = 2000 # steps per rollout episode (double well) + summarize_only: bool = False # skip the runs and summarize the files already in output_dir + plot: bool = False # also draw the two figures of `plot` into output_dir + + +# -------------------------------------------------------------------------------------------------------------------- +# Training and the deployed policy +# -------------------------------------------------------------------------------------------------------------------- + + +def load_config(env_id): + """Tuned SKVI configuration for `env_id`, from its file in `configurations/`.""" + with open(os.path.join(REPO_ROOT, "configurations", CONFIG_FILES[env_id])) as f: + return json.load(f) + + +def quadratic_cost(env, action_cost_scale=1.0): + """The environment's cost (x - x_ref)^T Q (x - x_ref) + s u^T R u as an (actions, states) matrix. + + With the default s = 1 this is the benchmark cost, with the same values as `env.vectorized_cost_fn`, which forms + an N x N intermediate and runs out of memory at the SKVI batch size of 2^14 states. R is diagonal in every + environment of the package. `skvi_sensitivity_checks` also trains with s = dt, which makes actions cheaper relative + to the state cost by that factor. + """ + Q = torch.as_tensor(env.unwrapped.Q, dtype=torch.float64) + R_diagonal = torch.diag(torch.as_tensor(env.unwrapped.R, dtype=torch.float64)) * action_cost_scale + reference = torch.as_tensor(env.unwrapped.reference_point, dtype=torch.float64) + + def cost(states, actions): + deviation = torch.as_tensor(np.asarray(states), dtype=torch.float64) - reference # (N, d) + state_cost = torch.einsum("ni,ij,nj->n", deviation, Q, deviation) # (N,) + action_cost = torch.as_tensor(np.asarray(actions), dtype=torch.float64) ** 2 @ R_diagonal # (L,) + return (state_cost.unsqueeze(-1) + action_cost.unsqueeze(0)).T # (L, N) + + return cost + + +def identify_tensor(env_id, seed, config, dt, extra_transitions=None): + """Identify the Koopman tensor from random-agent data as the SKVI script does, optionally with more transitions. + + extra_transitions (X, U, Y) arrays of shape (n, d), (n, 1), (n, d) appended to the random-agent data before the + tensor is refitted, for example transitions of a trained policy (re-identification along the + policy). SKVI then also fits its value function on these states, because it fits it on the + identification states. They are rounded to single precision before they are appended (a + relative change below 1e-7, far below the model's one-step error), which keeps the results + identical to those of the scripts these checks were first run with. + """ + tensor = generate_koopman_tensor( + env_id=env_id, + seed=seed, + num_paths=config["num-paths"], + num_steps_per_path=config["num-steps-per-path"], + state_order=config["state-order"], + action_order=config["action-order"], + regressor="ols", + ) + if extra_transitions is None: + return tensor + X_extra, U_extra, Y_extra = (torch.from_numpy(a.T.astype(np.float32)) for a in extra_transitions) + return KoopmanTensor( + torch.cat([tensor.X, X_extra], 1), + torch.cat([tensor.Y, Y_extra], 1), + torch.cat([tensor.U, U_extra], 1), + phi=monomial_dictionary(config["state-order"]), + psi=monomial_dictionary(config["action-order"]), + regressor=Regressor("ols"), + **({} if dt is None else {"dt": dt}), + ) + + +def train_skvi( + env_id, seed, state_order=None, action_cost_scale=1.0, extra_transitions=None, extra_transitions_for="both" +): + """Identify the Koopman tensor and train SKVI exactly as the SKVI script does, without writing checkpoints. + + The defaults give the tuned configuration in `configurations/`. The optional arguments, used by + `skvi_sensitivity_checks`, change the order of the state dictionary, scale the action-cost weight R (in training + and in the deployed policy) and append transitions to the identification data (see `identify_tensor`). SKVI draws + the states on which it fits its value function from the identification states, so appended transitions change + both the tensor and those states. `extra_transitions_for` separates the two: "both" (re-identification), "tensor" + (the value function is fitted on the random-agent states only) or "value" (the tensor is fitted on the + random-agent data only and the appended states join the value fit). + """ + if extra_transitions_for not in ("both", "tensor", "value"): + raise ValueError("extra_transitions_for must be 'both', 'tensor' or 'value'") + config = load_config(env_id) + if state_order is not None: + config = dict(config, **{"state-order": state_order}) + np.random.seed(seed) + torch.manual_seed(seed) + env = gym.make(env_id) + dt = getattr(env.unwrapped, "dt", None) + with contextlib.redirect_stdout(io.StringIO()): # the package prints progress for every epoch + for_tensor = extra_transitions if extra_transitions_for in ("both", "tensor") else None + tensor = identify_tensor(env_id, seed, config, dt, for_tensor) + if extra_transitions is not None and extra_transitions_for != "both": + random_agent = config["num-paths"] * config["num-steps-per-path"] + if extra_transitions_for == "tensor": # fit the value function on the random-agent states only + tensor.X, tensor.Phi_X = tensor.X[:, :random_agent], tensor.Phi_X[:, :random_agent] + else: # add the transitions' states to the value fit, keeping the random-agent tensor + states = torch.from_numpy(extra_transitions[0].T.astype(np.float32)).to(tensor.X.dtype) + tensor.X = torch.cat([tensor.X, states], 1) + tensor.Phi_X = torch.cat([tensor.Phi_X, tensor.phi(states).to(tensor.Phi_X.dtype)], 1) + actions = torch.from_numpy(np.linspace(env.action_space.low, env.action_space.high, NUM_ACTIONS)).T + policy = DiscreteKoopmanValueIterationPolicy( + env_id=env_id, + gamma=GAMMA, + alpha=ALPHA, + dynamics_model=tensor, + all_actions=actions, + cost=quadratic_cost(env, action_cost_scale), + seed=seed, + use_ols=True, + dt=dt, + ) + policy.train(config["number-of-train-epochs"], BATCH_SIZE, 1, how_often_to_chkpt=10**9) + return env, tensor, policy, config + + +def gibbs_policy(policy, states): + """Action probabilities (actions, states) of the deployed SKVI policy at `states` (N, d). + + Identical to `DiscreteKoopmanValueIterationPolicy.pis`, which accepts one state at a time: the score of action u + is -(c(x, u) + discount * w^T K^u phi(x)) / alpha, shifted by its maximum, exponentiated, floored by machine + epsilon and normalised over the action grid. + """ + states = np.asarray(states, dtype=np.float64) + tensor = policy.dynamics_model + phi = tensor.phi(torch.from_numpy(states.T).to(tensor.K.dtype)) # (d_phi, N) + K_u = tensor.K_(policy.all_actions).to(phi.dtype) # (L, d_phi, d_phi) + w = policy.value_function_weights.reshape(-1).to(phi.dtype) # (d_phi,) + continuation = torch.einsum("p,apq,qn->an", w, K_u, phi) # (L, N) + cost = torch.as_tensor(policy.cost(torch.from_numpy(states), policy.all_actions.T), dtype=continuation.dtype) + score = -(cost + policy.discount_factor * continuation) / policy.alpha + unnormalised = torch.exp(score - score.amax(dim=0, keepdim=True)) + PROBABILITY_FLOOR + return (unnormalised / unnormalised.sum(dim=0, keepdim=True)).detach().numpy() + + +def mean_action(policy, states): + """Mean action of the deployed policy at `states` (N, d), for a scalar action.""" + grid = policy.all_actions.numpy().reshape(-1, 1) + return (grid * gibbs_policy(policy, states)).sum(axis=0) + + +# -------------------------------------------------------------------------------------------------------------------- +# Reading the value function and the action fields +# -------------------------------------------------------------------------------------------------------------------- + + +def monomials(state_dim, order): + """Names and exponent matrix (state_dim, d_phi) of the package's monomial dictionary.""" + powers = allMonomialPowers(state_dim, order).numpy().astype(int) + names = [] + for k in range(powers.shape[1]): + factors = [f"x{i + 1}" if p == 1 else f"x{i + 1}^{p}" for i, p in enumerate(powers[:, k]) if p > 0] + names.append("".join(factors) or "1") + return names, powers + + +def quadratic_form(coefficients, powers): + """Symmetric matrix P with x^T P x equal to the degree-two part of sum_k coefficients[k] * monomial_k(x).""" + P = np.zeros((powers.shape[0],) * 2) + for k in range(powers.shape[1]): + if powers[:, k].sum() != 2: + continue + variables = np.nonzero(powers[:, k])[0] + if len(variables) == 1: # x_i^2 + i = variables[0] + P[i, i] = coefficients[k] + else: # x_i x_j, split evenly between the two off-diagonal entries + i, j = variables + P[i, j] = P[j, i] = coefficients[k] / 2 + return P + + +def linear_coefficients(coefficients, powers): + """Coefficients of the degree-one monomials, in state order.""" + linear = np.zeros(powers.shape[0]) + for k in range(powers.shape[1]): + if powers[:, k].sum() == 1: + linear[np.argmax(powers[:, k])] = coefficients[k] + return linear + + +def action_fields(tensor, w): + """Coefficient vectors of the fields h_k(x) = w^T T[:, :, k] phi(x), one row per action-dictionary function. + + With the polynomial action dictionary psi(u) = (1, u, u^2, ...), the fitted continuation is a polynomial in the + action, w^T K^u phi(x) = sum_k h_k(x) u^k, and row k holds the coefficients of h_k on the state dictionary phi. + The deployed policy weighs action u by exp(-(c(x, u) + discount * sum_k h_k(x) u^k) / alpha), so it is an + exponential family over the action dictionary with natural parameters -discount * h_k(x) / alpha and depends on w + only through these fields. This is Proposition S11.1 ("The policy is an exponential family over the action + dictionary") of the reference in the module docstring. The reference states it for the value function, whose + fields have the opposite sign, because the package stores the coefficients w of the cost-to-go. + """ + T = tensor.K.numpy() + return np.array([T[:, :, k].T @ w for k in range(T.shape[2])]) + + +def discounted_lqr(A, B, Q, R, gamma): + """Riccati matrix P and gain K (u = -K x) of the discounted LQR.""" + A_discounted, B_discounted = np.sqrt(gamma) * A, np.sqrt(gamma) * B + P = solve_discrete_are(A_discounted, B_discounted, Q, R) + K = np.linalg.solve(R + B_discounted.T @ P @ B_discounted, B_discounted.T @ P @ A_discounted) + return P, K + + +def relative_error(estimate, reference): + """Ratio of Euclidean (Frobenius for matrices) norms.""" + return float(np.linalg.norm(np.asarray(estimate) - np.asarray(reference)) / np.linalg.norm(reference)) + + +# -------------------------------------------------------------------------------------------------------------------- +# The checks +# -------------------------------------------------------------------------------------------------------------------- + + +def policy_reading(env_id, env, tensor, policy, w, names, powers): + """Gain of the mean action against an LQR comparator, and the closed-form Gaussian mean. + + When the action cost is R u^2 and the fields h_k of degree three and higher vanish (`action_fields`), the policy's + weights on the action grid follow a Gaussian with mean -discount * h_1(x) / (2 (R + discount * h_2(x))) and + variance alpha / (2 (R + discount * h_2(x))). This is Corollary S11.2 ("Polynomial action dictionary, quadratic + action cost") of the reference in the module docstring, with the sign of the fields as in `action_fields`. With + h_2 constant, the degree-one part of h_1 gives the closed-form gain, which is compared with the gain fitted to the + grid policy's mean action. On the double well the closed-form mean is also compared with the grid mean on a box + of states. + + Linear system: the comparator is the discounted LQR of the seed's (A, B, Q, R), on states whose LQR action lies + well inside the action grid (|Kx| < 7 for the grid [-10, 10]). Double well: the comparator is the LQR of the + drift linearised at the origin, on states near the origin. + """ + state_dim = powers.shape[0] + discount = policy.discount_factor + R = float(env.unwrapped.R[0, 0]) + fields = action_fields(tensor, w) + constant_index = names.index("1") + h2_constant = fields[2][constant_index] if fields.shape[0] > 2 else 0.0 + closed_form_gain = discount * linear_coefficients(fields[1], powers) / (2 * (R + discount * h2_constant)) + + if env_id == "LinearSystem-v0": + A, B = env.unwrapped.A, env.unwrapped.B + P, K = discounted_lqr(A, B, env.unwrapped.Q, env.unwrapped.R, discount) + states = np.random.uniform(-25, 25, size=(6000, state_dim)) + states = states[np.abs(states @ K.T).reshape(-1) < 7.0] + else: + A = np.eye(2) + env.unwrapped.dt * np.diag([4.0, -2.0]) + B = env.unwrapped.dt * np.ones((2, 1)) + P, K = discounted_lqr(A, B, np.eye(2), np.eye(1), discount) + states = np.random.uniform(-0.3, 0.3, size=(6000, state_dim)) + grid_gain = -np.linalg.lstsq(states, mean_action(policy, states), rcond=None)[0] + learned_P = quadratic_form(w, powers) + + reading = { + "fields": fields.tolist(), + "h2_nonconstant_share": ( + float(np.abs(np.delete(fields[2], constant_index)).sum() / max(np.abs(fields[2]).sum(), 1e-12)) + if fields.shape[0] > 2 + else None + ), + "h3_relative": ( + float(np.abs(fields[3]).max() / max(np.abs(fields[1]).max(), 1e-12)) if fields.shape[0] > 3 else None + ), + "P_lqr": P.tolist(), + "P_learned": learned_P.tolist(), + "P_relative_error": relative_error(learned_P, P), + "K_lqr": K.reshape(-1).tolist(), + "K_closed_form": closed_form_gain.tolist(), + "K_grid": grid_gain.tolist(), + "K_grid_relative_error": relative_error(grid_gain, K.reshape(-1)), + "K_closed_form_vs_grid": relative_error(closed_form_gain, grid_gain), + "discount": discount, + "n_states": int(len(states)), + } + if env_id == "DoubleWell-v0": + axis = np.linspace(-2, 2, 41) + X1, X2 = np.meshgrid(axis, axis) + points = np.c_[X1.ravel(), X2.ravel()] + grid_mean = mean_action(policy, points) + h = [sum(c * np.prod(points ** powers[:, k], axis=1) for k, c in enumerate(fields[z])) for z in (1, 2)] + closed_form_mean = -discount * h[0] / (2 * (R + discount * h[1])) + reading["mean_action_field"] = {"axis": axis.tolist(), "grid_mean": grid_mean.reshape(41, 41).tolist()} + reading["closed_form_vs_grid_max_abs"] = float(np.abs(closed_form_mean - grid_mean).max()) + return reading + + +def rollout_returns(env, policy, seed, episodes, horizon, record_every=None): + """Returns of `episodes` rollouts of the deployed policy, with common random numbers across calls. + + Episode e is seeded with seed * 1000 + e, so pruned and full policies meet the same initial states and noise. + With `record_every`, the states visited every that many steps are returned as well. + """ + returns, visited = [], [] + for episode in range(episodes): + np.random.seed(seed * 1000 + episode) + observation = env.reset() + total = 0.0 + for step in range(horizon): + if record_every and step % record_every == 0: + visited.append(np.asarray(observation, dtype=np.float64).copy()) + probabilities = gibbs_policy(policy, np.asarray(observation, dtype=np.float64).reshape(1, -1))[:, 0] + action = policy.all_actions[0][np.random.choice(len(probabilities), p=probabilities)].reshape(-1).numpy() + observation, reward, done, _ = env.step(action) + total += float(reward) + if done: + break + returns.append(total) + return returns, np.array(visited) + + +def pruning_experiment(env, tensor, policy, w, powers, seed, episodes, horizon): + """Double well: does the policy need the small terms of the learned value function? + + Variants: the full value function; the odd terms (x, y) removed; only the constant and the squares kept. Returns + are compared on paired episodes, and the total-variation distance to the full policy is measured on the states + the full policy visits. + """ + states = np.random.uniform(-2, 2, size=(20000, 2)) + phi = tensor.phi(torch.from_numpy(states.T)).numpy() + degree = powers.sum(axis=0) + variables = (powers > 0).sum(axis=0) + masks = { + "even": (degree % 2 == 0).astype(float), # drop the odd monomials x and y + "diag": ((degree % 2 == 0) & (variables <= 1)).astype(float), # keep 1, x^2 and y^2 + } + out = {"importance": (np.abs(w) * phi.std(axis=1)).tolist()} + + policy.value_function_weights = torch.from_numpy(w.reshape(-1, 1)) + out["returns_full"], visited = rollout_returns(env, policy, seed, episodes, horizon, record_every=20) + full = gibbs_policy(policy, visited) + for variant in ("even", "diag"): + policy.value_function_weights = torch.from_numpy((w * masks[variant]).reshape(-1, 1)) + tv = 0.5 * np.abs(gibbs_policy(policy, visited) - full).sum(axis=0) + out[f"tv_{variant}"] = {"mean": float(tv.mean()), "max": float(tv.max())} + out[f"returns_{variant}"] = rollout_returns(env, policy, seed, episodes, horizon)[0] + policy.value_function_weights = torch.from_numpy(w.reshape(-1, 1)) + return out + + +def run(env_id, seed, args): + """Train SKVI for one seed, run the checks and write the result file.""" + start = time.time() + env, tensor, policy, config = train_skvi(env_id, seed) + names, powers = monomials(env.observation_space.shape[0], config["state-order"]) + w = policy.value_function_weights.detach().clone().flatten().numpy() + result = {"env": env_id, "seed": seed, "names": names, "w": w.tolist(), "train_seconds": time.time() - start} + result["reading"] = policy_reading(env_id, env, tensor, policy, w, names, powers) + if env_id == "DoubleWell-v0": + result.update(pruning_experiment(env, tensor, policy, w, powers, seed, args.episodes, args.horizon)) + result["total_seconds"] = time.time() - start + path = os.path.join(args.output_dir, f"{env_id.split('-')[0]}_{seed}.json") + with open(path, "w") as f: + json.dump(result, f, indent=1) + reading = result["reading"] + if env_id == "LinearSystem-v0": + detail = ( + f"|P_learned - P|/|P| = {reading['P_relative_error']:.3f}, " + f"|K_grid - K|/|K| = {reading['K_grid_relative_error']:.3f}" + ) + else: + full = np.mean(result["returns_full"]) + changes = [100 * (np.mean(result[f"returns_{v}"]) - full) / abs(full) for v in ("even", "diag")] + detail = ( + f"gain of the mean action {np.round(reading['K_grid'], 3).tolist()}, " + f"return change when pruned {changes[0]:+.2f}% / {changes[1]:+.2f}%" + ) + print(f"{env_id} seed {seed}: {detail} ({result['total_seconds'] / 60:.1f} min)", flush=True) + + +# -------------------------------------------------------------------------------------------------------------------- +# Summary and figures +# -------------------------------------------------------------------------------------------------------------------- + + +def median_iqr(values, digits=3): + values = np.asarray(values, dtype=float) + q1, median, q3 = np.percentile(values, [25, 50, 75]) + return f"{median:.{digits}f} [{q1:.{digits}f}, {q3:.{digits}f}]" + + +def load_results(output_dir, name): + return [json.load(open(p)) for p in sorted(glob.glob(os.path.join(output_dir, f"{name}_*.json")))] + + +def summarize(output_dir): + """Print the results of all seeds in output_dir: medians and interquartile ranges over seeds, and the extremes.""" + linear = load_results(output_dir, "LinearSystem") + if linear: + P_err = [r["reading"]["P_relative_error"] for r in linear] + K_err = [r["reading"]["K_grid_relative_error"] for r in linear] + agreement = [r["reading"]["K_closed_form_vs_grid"] for r in linear] + print(f"Linear system, {len(linear)} seeds") + print(f" |P_learned - P|/|P| median [IQR] {median_iqr(P_err)}, max {max(P_err):.3f}") + print(f" |K_grid - K|/|K| median [IQR] {median_iqr(K_err)}, max {max(K_err):.3f}") + print(f" closed form vs grid gain median [IQR] {median_iqr(agreement, 5)}") + print(f" h2 non-constant share max {max(r['reading']['h2_nonconstant_share'] for r in linear):.2e}") + print(f" |h3| / |h1| max {max(r['reading']['h3_relative'] for r in linear):.2e}") + + well = load_results(output_dir, "DoubleWell") + if well: + names = well[0]["names"] + W = np.array([r["w"] for r in well]) + importance = np.array([r["importance"] for r in well]) + share = importance / importance.sum(axis=1, keepdims=True) + odd_share = share[:, [names.index("x1"), names.index("x2")]].sum(axis=1) + print(f"Double well, {len(well)} seeds (cost-to-go coefficients as stored by the package)") + for j, name in enumerate(names): + share_percent = 100 * np.median(share[:, j]) + print(f" {name:>5}: coefficient {median_iqr(W[:, j], 2)}, importance share {share_percent:.1f}%") + print(f" odd terms together: median {100 * np.median(odd_share):.1f}%, max {100 * odd_share.max():.1f}%") + full = np.array([np.mean(r["returns_full"]) for r in well]) + print(f" return, full value function: {median_iqr(full, 1)}") + for variant in ("even", "diag"): + pruned = np.array([np.mean(r[f"returns_{variant}"]) for r in well]) + change = 100 * (pruned - full) / np.abs(full) + tv_max = max(r[f"tv_{variant}"]["max"] for r in well) + print( + f" {variant:>4}: relative change of the return {median_iqr(change, 2)} % " + f"(min {change.min():+.2f}, max {change.max():+.2f}), largest TV distance {tv_max:.3f}" + ) + gains = np.array([r["reading"]["K_grid"] for r in well]) + agreement = [r["reading"]["closed_form_vs_grid_max_abs"] for r in well] + print(f" gain of the mean action near the origin: {np.round(np.median(gains, axis=0), 3).tolist()}") + print(f" closed-form vs grid mean action, largest |difference| on the box: median {np.median(agreement):.1e}") + + +def plot(output_dir): + """Draw two figures into output_dir. + + linear_costate.pdf linear system, all seeds: gain from the mean action against the LQR gain, and the + learned quadratic form against the Riccati matrix + double_well_costate.pdf double well, first seed: the learned cost-to-go and the mean action on [-2, 2]^2, with + the direction of the drift + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + linear = load_results(output_dir, "LinearSystem") + if linear: + fig, ax = plt.subplots(1, 2, figsize=(7.6, 3.3)) + K = np.concatenate([r["reading"]["K_lqr"] for r in linear]) + K_grid = np.concatenate([r["reading"]["K_grid"] for r in linear]) + upper = np.triu_indices(3) + P = np.concatenate([np.array(r["reading"]["P_lqr"])[upper] for r in linear]) + P_learned = np.concatenate([np.array(r["reading"]["P_learned"])[upper] for r in linear]) + for axis, (x, y, xlabel, ylabel) in zip( + ax, + [ + (K, K_grid, "LQR gain component $K_i$", "gain from the Gibbs mean action"), + (P, P_learned, "Riccati entry $P_{ij}$", "learned quadratic form"), + ], + ): + axis.scatter(x, y, s=18) + low, high = min(x.min(), y.min(), 0) * 1.05, max(x.max(), y.max()) * 1.05 + axis.plot([low, high], [low, high], "k--", lw=0.8) + axis.set_xlabel(xlabel) + axis.set_ylabel(ylabel) + fig.tight_layout() + fig.savefig(os.path.join(output_dir, "linear_costate.pdf")) + + well = load_results(output_dir, "DoubleWell") + if well: + result = well[0] + field = result["reading"]["mean_action_field"] + axis_values = np.array(field["axis"]) + X1, X2 = np.meshgrid(axis_values, axis_values) + points = np.c_[X1.ravel(), X2.ravel()] + _, powers = monomials(2, 2) + V = sum(c * np.prod(points ** powers[:, k], axis=1) for k, c in enumerate(result["w"])).reshape(X1.shape) + U = np.array(field["grid_mean"]) + fig, ax = plt.subplots(1, 2, figsize=(8.2, 3.4)) + fig.colorbar(ax[0].contourf(X1, X2, V, 30, cmap="viridis"), ax=ax[0]) + ax[0].set_title("learned soft cost-to-go $\\hat V(x)$") + largest = np.abs(U).max() + fig.colorbar(ax[1].contourf(X1, X2, U, 30, cmap="RdBu_r", vmin=-largest, vmax=largest), ax=ax[1]) + ax[1].set_title("Gibbs mean action (arrows: drift direction)") + coarse = axis_values[::5] + A1, A2 = np.meshgrid(coarse, coarse) + F1, F2 = 4 * A1 - 4 * A1**3, -2 * A2 + norm = np.sqrt(F1**2 + F2**2) + 1e-9 + ax[1].quiver(A1, A2, F1 / norm, F2 / norm, color="k", alpha=0.45, scale=28, width=0.004, headwidth=4) + for axis in ax: + axis.set_xlabel("$x$") + axis.set_ylabel("$y$") + fig.tight_layout() + fig.savefig(os.path.join(output_dir, "double_well_costate.pdf")) + + +if __name__ == "__main__": + args = ArgumentParser().parse_args() + os.makedirs(args.output_dir, exist_ok=True) + torch.set_num_threads(4) + if not args.summarize_only: + for environment in args.environments: + for seed in args.seeds: + run(f"{environment}-v0", seed, args) + summarize(args.output_dir) + if args.plot: + plot(args.output_dir) diff --git a/koopmanrl_utils/skvi_sensitivity_checks.py b/koopmanrl_utils/skvi_sensitivity_checks.py new file mode 100644 index 0000000..16c950d --- /dev/null +++ b/koopmanrl_utils/skvi_sensitivity_checks.py @@ -0,0 +1,676 @@ +"""Accuracy of the fitted Koopman tensor along the trained SKVI policy, and sensitivity of SKVI's control. + +The Koopman tensor is identified from random-agent data, as in `koopmanrl.soft_koopman_value_iteration`, while the +trained policy visits other states. The first study measures how accurate the fitted model is along the policy. The +second retrains SKVI with other dictionaries, action costs and identification data, and compares its control with +references that do not use the tensor. + +Accuracy along the policy (`--study accuracy`; all four benchmarks) + The tensor is evaluated on transitions of the trained SKVI policy and on fresh random-agent transitions. The + one-step error is the root mean square, over transitions n and nonconstant dictionary functions j, of the + normalised prediction error + + (K^{u_n} phi(x_n) - E[phi(x'_n) | x_n, u_n])_j / s_j, + + where the conditional mean is phi(x'_n) itself for the deterministic benchmarks and the exact conditional mean of + the quadratic dictionary under one Euler-Maruyama step for the double well. The scale s_j is the root mean square + of the target on the random-agent transitions and is used for both distributions ("common scale"), so that the two + errors are comparable. Recomputing it on the data being evaluated ("own scale") inflates the error near the + target, where the high-degree functions are tiny, by orders of magnitude, so that version is recorded too. The same + error with phi(x_n) in place of the prediction is the persistence baseline (predicting no change). The error of + the predicted state is reported in the units of the state. The run also records the one-step Jacobian of the + fitted model at the target against that of the environment's one-step map, and the distance from the target of the + closed loop and of the uncontrolled system. + +Sensitivity of the control (`--study sensitivity`; double well and Lorenz) + SKVI is retrained with other orders of the state dictionary, with the action-cost weight R multiplied by the time + step (actions a hundred times cheaper relative to the state cost), and after re-identifying the tensor on + transitions of the trained policy. Each variant is compared with no control and with the LQR controller of the + linearisation at the target, on the same episodes. Only SKVI is retrained. + +Usage (from the repository root): + + uv run -m koopmanrl_utils.skvi_sensitivity_checks --study accuracy # four benchmarks, seeds 100-107 + uv run -m koopmanrl_utils.skvi_sensitivity_checks --study sensitivity # double well and Lorenz + uv run -m koopmanrl_utils.skvi_sensitivity_checks --study sensitivity --environments Lorenz --seeds 100 101 + uv run -m koopmanrl_utils.skvi_sensitivity_checks --summarize_only + +Each run writes one JSON file into `--output_dir` (default `skvi_sensitivity_checks_results/`, not tracked by git), +with the settings it used; the summary refuses to combine runs made with different settings. Random draws use NumPy's +global generator, seeded per episode or per data set in a fixed order, so a run is reproducible from its seed. Tensors +and policies are built by `skvi_policy_checks.train_skvi`, as in `koopmanrl.soft_koopman_value_iteration`, with the +tuned configurations in `configurations/`. + +The results were first produced by exploratory scripts that this module consolidates. With the same seeds it +reproduces their numbers exactly, except the one-step Jacobians of the fitted model recorded by the sensitivity study, +which it evaluates in double precision with a step of 1e-4 (differences below 1e-7). + +Reference: the results of these checks are in the electronic supplementary material of "Koopman-Assisted Reinforcement +Learning" (Rozwood, Mehrez, Paehler, Sun and Brunton), section "Accuracy along the learned policy and sensitivity of +SKVI". +""" + +import contextlib +import glob +import io +import json +import os +import time + +import gym +import numpy as np +import torch +from scipy.linalg import solve_discrete_are +from tap import Tap + +import koopmanrl.environments # noqa: F401 (registers the environments with gym) +from koopmanrl_utils.skvi_policy_checks import PROBABILITY_FLOOR, monomials, train_skvi + +ACCURACY_ENVIRONMENTS = ["LinearSystem", "DoubleWell", "FluidFlow", "Lorenz"] +SENSITIVITY_ENVIRONMENTS = ["DoubleWell", "Lorenz"] + +# Steps per closed-loop episode: the benchmark episode of 2,000 steps, and 200 for the discrete-time linear system. +HORIZON = {"LinearSystem-v0": 200, "DoubleWell-v0": 2000, "FluidFlow-v0": 2000, "Lorenz-v0": 2000} + +# A state is near the target when its Euclidean distance from the target is below this radius; the sensitivity study +# reports the fraction of each episode spent near the target. +NEAR_RADIUS = {"DoubleWell-v0": 0.25, "Lorenz-v0": 2.0} +# A seed counts as failed when its mean-action policy is near the target for less than this fraction of the episode. +FAILURE_FRACTION = 0.001 + +# Variants of the sensitivity study: (name, state-dictionary order, action-cost scale, rounds of re-identification, +# what the policy transitions are used for: "both" re-identifies, "tensor" or "value" separates the two effects, see +# `skvi_policy_checks.train_skvi`). The action-cost scale 0.01 is R * dt for both systems (dt = 0.01). +VARIANTS = { + "DoubleWell-v0": [ + ("order 2", 2, 1.0, 0, "both"), + ("order 2, R*dt", 2, 0.01, 0, "both"), + ("order 4, R*dt", 4, 0.01, 0, "both"), + ("order 5, R*dt", 5, 0.01, 0, "both"), + ("order 6, R*dt", 6, 0.01, 0, "both"), + ], + "Lorenz-v0": [ + ("order 3", 3, 1.0, 0, "both"), + ("order 3, refit x1", 3, 1.0, 1, "both"), + ("order 3, refit x1, tensor only", 3, 1.0, 1, "tensor"), + ("order 3, refit x1, value states only", 3, 1.0, 1, "value"), + ("order 4", 4, 1.0, 0, "both"), + ("order 4, refit x1", 4, 1.0, 1, "both"), + ("order 4, refit x2", 4, 1.0, 2, "both"), + ], +} +# LQR references of the sensitivity study: (name, action-cost scale). +LQR_REFERENCES = { + "DoubleWell-v0": [("LQR, benchmark R", 1.0), ("LQR, R*dt", 0.01)], + "Lorenz-v0": [("LQR, benchmark R", 1.0)], +} + + +class ArgumentParser(Tap): + study: str = "accuracy" # "accuracy" (model error along the policy) or "sensitivity" (SKVI retrained in variants) + environments: list[str] = [] # default: all four benchmarks for accuracy, double well and Lorenz for sensitivity + seeds: list[int] = list(range(100, 108)) # one SKVI training per seed and variant + output_dir: str = "skvi_sensitivity_checks_results" # one JSON file per study, environment and seed + accuracy_episodes: int = 16 # closed-loop episodes sampled for the accuracy study + random_agent_paths: int = 20 # random-agent trajectories for the reference distribution + evaluation_episodes: int = 10 # closed-loop episodes per controller in the sensitivity study + variants: list[str] = [] # sensitivity study: run only these variants (default: all); each variant is seeded alone + summarize_only: bool = False # skip the runs and summarize the files already in output_dir + + +# -------------------------------------------------------------------------------------------------------------------- +# Training and the deployed policy +# -------------------------------------------------------------------------------------------------------------------- + + +def gibbs_sampler(policy): + """Fast deployed policy for one state at a time: (probabilities(x), action grid). + + probabilities(x) returns the Gibbs weights over the action grid, computed as in the package (score + -(c(x, u) + discount * w^T K^u phi(x)) / alpha, shifted by its maximum, exponentiated, floored by machine epsilon + and normalised), in double precision with w^T K^u precomputed for every grid action. + """ + tensor = policy.dynamics_model + K_u = tensor.K_(policy.all_actions).to(torch.float64) # (L, d_phi, d_phi) + w = policy.value_function_weights.reshape(-1).to(torch.float64) + continuation = torch.einsum("p,apq->aq", w, K_u).numpy() # (L, d_phi) + grid = policy.all_actions.numpy().reshape(-1) + + def probabilities(x): + x = np.asarray(x, dtype=np.float64) + phi = tensor.phi(torch.from_numpy(x.reshape(-1, 1)).to(tensor.K.dtype)).numpy().astype(np.float64).reshape(-1) + cost = policy.cost(x.reshape(1, -1), grid.reshape(-1, 1)).numpy().reshape(-1) + score = -(cost + policy.discount_factor * (continuation @ phi)) / policy.alpha + unnormalised = np.exp(score - score.max()) + PROBABILITY_FLOOR + return unnormalised / unnormalised.sum() + + return probabilities, grid + + +def deployed_controller(policy, mean=False): + """Controller x -> action of the deployed policy: a sampled grid action, or with mean=True the mean action.""" + probabilities, grid = gibbs_sampler(policy) + + def act(x): + p = probabilities(x) + if mean: + return np.array([float(p @ grid)]) + return np.array([grid[np.random.choice(len(p), p=p)]]) + + return act + + +# -------------------------------------------------------------------------------------------------------------------- +# Episodes and data +# -------------------------------------------------------------------------------------------------------------------- + + +def closed_loop_episode(env, controller, seed, horizon): + """States, actions and next states of one episode from the environment's initial state, seeded with `seed`.""" + np.random.seed(seed) + observation = env.reset() + X, U, Y = [], [], [] + for _ in range(horizon): + x = np.asarray(observation, dtype=np.float64).copy() + u = controller(x) + observation, _, done, _ = env.step(u) + X.append(x), U.append(float(u[0])), Y.append(np.asarray(observation, dtype=np.float64).copy()) + if done: + break + return np.array(X), np.array(U), np.array(Y) + + +def random_agent_data(env, seed, paths, steps): + """Transitions of the uniform random agent that identifies the tensor, from fresh trajectories.""" + np.random.seed(seed) + env.action_space.seed(seed) + X, U, Y = [], [], [] + for _ in range(paths): + observation = env.reset() + for _ in range(steps): + x = np.asarray(observation, dtype=np.float64).copy() + u = env.action_space.sample() + observation, _, _, _ = env.step(u) + ( + X.append(x), + U.append(float(np.asarray(u).reshape(-1)[0])), + Y.append(np.asarray(observation, dtype=np.float64).copy()), + ) + return np.array(X), np.array(U), np.array(Y) + + +def policy_transitions(env, policy, paths, steps, seed): + """Transitions of the trained (sampled) policy from the environment's initial states, for re-identification.""" + act = deployed_controller(policy) + np.random.seed(seed) + X, U, Y = [], [], [] + for _ in range(paths): + observation = env.reset() + for _ in range(steps): + x = np.asarray(observation, dtype=np.float64) + u = act(x) + observation, _, done, _ = env.step(u) + X.append(x), U.append(u), Y.append(np.asarray(observation, dtype=np.float64)) + if done: + break + return np.array(X), np.array(U), np.array(Y) + + +# -------------------------------------------------------------------------------------------------------------------- +# One-step accuracy +# -------------------------------------------------------------------------------------------------------------------- + + +def features(tensor, states): + """Dictionary values (N, d_phi) at `states` (N, d), in double precision.""" + states = np.asarray(states, dtype=np.float64) + return tensor.phi(torch.from_numpy(states.T).to(tensor.K.dtype)).numpy().astype(np.float64).T + + +def predict(tensor, Phi, U): + """Fitted one-step prediction K^u phi(x) per transition, with the action dictionary psi(u) = (1, u, u^2, ...).""" + T = tensor.K.numpy().astype(np.float64) # (d_phi, d_phi, d_psi) + Psi = np.stack([np.asarray(U, dtype=np.float64) ** k for k in range(T.shape[2])], axis=1) + return np.einsum("ijk,nj,nk->ni", T, Phi, Psi) + + +def linear_indices(powers): + """Positions of the degree-one monomials x_1, ..., x_d in the dictionary.""" + d = powers.shape[0] + return [int(np.where((powers == np.eye(d, dtype=int)[:, [i]]).all(axis=0))[0][0]) for i in range(d)] + + +def double_well_conditional_mean(env, X, U, powers): + """Exact E[phi(x') | x, u] and E[x' | x, u] for a dictionary of degree at most two, under one Euler-Maruyama step. + + x' = x + f(x, u) dt + S(x) xi sqrt(dt) with S(x) = [[0.7, x_1], [0, 0.5]] (the environment's noise), so + E[x'] = m and E[x' x'^T] = m m^T + S S^T dt. + """ + if powers.sum(axis=0).max() > 2: + raise ValueError("the closed form holds for dictionaries of degree at most two") + dt = env.unwrapped.dt + targets, means = [], [] + for x, u in zip(X, U): + m = x + np.asarray(env.unwrapped.continuous_f(np.array([u]))(0, x), dtype=np.float64).reshape(-1) * dt + S = np.array([[0.7, x[0]], [0.0, 0.5]]) + C = S @ S.T * dt + row = [] + for p in powers.T: + variables = np.nonzero(p)[0] + if p.sum() == 0: + row.append(1.0) + elif p.sum() == 1: + row.append(m[variables[0]]) + else: # x_i^2 or x_i x_j + i, j = variables[0], variables[-1] + row.append(m[i] * m[j] + C[i, j]) + targets.append(row), means.append(m) + return np.array(targets), np.array(means) + + +def one_step_errors(prediction, target, base, X, next_state, nonconstant, linear, common_scale): + """Normalised one-step errors of the model and of persistence, at two scales, and the error of the state. + + The normalised error is sqrt(mean_n mean_j ((prediction - target)_nj / s_j)^2) over the nonconstant dictionary + functions j, with s_j the root mean square of the target either on the data itself ("own scale") or on the + random-agent data ("common scale", which makes the two distributions comparable). The persistence error is the + same with `base` (phi(x), no change) in place of the prediction. + """ + p, t, b = prediction[:, nonconstant], target[:, nonconstant], base[:, nonconstant] + out = {"n": int(len(X))} + for name, scale in [("own_scale", np.sqrt(np.mean(t**2, axis=0))), ("common_scale", common_scale)]: + model = float(np.sqrt(np.mean(((p - t) / scale) ** 2))) + persistence = float(np.sqrt(np.mean(((b - t) / scale) ** 2))) + out[name] = {"model": model, "persistence": persistence, "ratio": model / persistence} + error = np.linalg.norm(prediction[:, linear] - next_state, axis=1) + change = np.linalg.norm(next_state - X, axis=1) + out["state"] = {"error_rms": float(np.sqrt(np.mean(error**2))), "change_rms": float(np.sqrt(np.mean(change**2)))} + return out + + +def one_step_map(env, x, u): + """The environment's one-step map without noise: the integrated flow, or the Euler drift for the double well.""" + if env.spec.id == "DoubleWell-v0": + drift = np.asarray(env.unwrapped.continuous_f(np.array([u]))(0, x), dtype=np.float64).reshape(-1) + return x + drift * env.unwrapped.dt + return np.asarray(env.unwrapped.f(np.asarray(x, dtype=np.float64), np.array([u])), dtype=np.float64).reshape(-1) + + +def local_dynamics(env, tensor, linear, h=1e-4): + """One-step Jacobians (in x and in u) at the target, of the fitted model and of the environment's one-step map.""" + target = np.asarray(env.unwrapped.reference_point, dtype=np.float64) + d = len(target) + + def model(x, u): + return predict(tensor, features(tensor, x.reshape(1, -1)), np.array([u]))[0, linear] + + def jacobians(step): + A = np.stack([(step(target + h * e, 0.0) - step(target - h * e, 0.0)) / (2 * h) for e in np.eye(d)], axis=1) + B = (step(target, h) - step(target, -h)) / (2 * h) + return A, B + + A_model, B_model = jacobians(model) + A_true, B_true = jacobians(lambda x, u: one_step_map(env, x, u)) + return { + "A_model": A_model.tolist(), + "A_true": A_true.tolist(), + "B_model": B_model.tolist(), + "B_true": B_true.tolist(), + "max_abs_entry_error": float(np.abs(A_model - A_true).max()), + "relative_error": float(np.linalg.norm(A_model - A_true) / np.linalg.norm(A_true - np.eye(d))), + "eigenvalue_moduli_model": sorted(np.abs(np.linalg.eigvals(A_model)).tolist()), + "eigenvalue_moduli_true": sorted(np.abs(np.linalg.eigvals(A_true)).tolist()), + } + + +def accuracy_run(env_id, seed, args): + """Accuracy study for one environment and seed: one-step error along the policy and on random-agent data.""" + start = time.time() + env, tensor, policy, config = train_skvi(env_id, seed) + d = env.observation_space.shape[0] + _, powers = monomials(d, config["state-order"]) + nonconstant = powers.sum(axis=0) > 0 + linear = linear_indices(powers) + target = np.asarray(env.unwrapped.reference_point, dtype=np.float64) + act = deployed_controller(policy) + horizon = HORIZON[env_id] + + # transitions of the policy: every few steps of each episode, so that every phase of the episode is represented + on_policy, distance, distance_free = [], [], [] + for episode in range(args.accuracy_episodes): + X, U, Y = closed_loop_episode(env, act, seed * 1000 + episode, horizon) + stride = max(1, len(X) // 250) + on_policy.append((X[::stride], U[::stride], Y[::stride])) + distance.append(np.linalg.norm(X - target, axis=1)) + X_free, _, _ = closed_loop_episode(env, lambda x: np.zeros(1), seed * 1000 + episode, horizon) + distance_free.append(np.linalg.norm(X_free - target, axis=1)) + X_on, U_on, Y_on = (np.concatenate(a) for a in zip(*on_policy)) + + X_r, U_r, Y_r = random_agent_data(env, seed + 50_000, args.random_agent_paths, config["num-steps-per-path"]) + keep = np.random.default_rng(seed).choice(len(X_r), min(len(X_r), 4000), replace=False) + X_r, U_r, Y_r = X_r[keep], U_r[keep], Y_r[keep] + + sets = {} + for name, (X, U, Y) in {"on_policy": (X_on, U_on, Y_on), "random_agent": (X_r, U_r, Y_r)}.items(): + Phi = features(tensor, X) + if env_id == "DoubleWell-v0": + target_features, next_state = double_well_conditional_mean(env, X, U, powers) + else: + target_features, next_state = features(tensor, Y), Y + sets[name] = (predict(tensor, Phi, U), target_features, Phi, X, next_state) + common_scale = np.sqrt(np.mean(sets["random_agent"][1][:, nonconstant] ** 2, axis=0)) + + result = { + "env": env_id, + "seed": seed, + "settings": { + "configuration": config, + "horizon": horizon, + "accuracy_episodes": args.accuracy_episodes, + "random_agent_paths": args.random_agent_paths, + }, + } + for name, (prediction, target_features, Phi, X, next_state) in sets.items(): + result[name] = one_step_errors( + prediction, target_features, Phi, X, next_state, nonconstant, linear, common_scale + ) + result["local_dynamics"] = local_dynamics(env, tensor, linear) + second_half = [r[len(r) // 2 :] for r in distance] + result["distance"] = { + "final_median": float(np.median([r[-1] for r in distance])), + "second_half_max": float(max(r.max() for r in second_half)), + "uncontrolled_final_median": float(np.median([r[-1] for r in distance_free])), + } + result["seconds"] = time.time() - start + write(args.output_dir, "accuracy", env_id, seed, result) + o, r = result["on_policy"]["common_scale"], result["random_agent"]["common_scale"] + print( + f"{env_id} seed {seed}: model error {o['model']:.3g} along the policy vs {r['model']:.3g} on random-agent data " + f"(persistence {o['persistence']:.3g} vs {r['persistence']:.3g}); local Jacobian relative error " + f"{result['local_dynamics']['relative_error']:.1e} ({result['seconds'] / 60:.1f} min)", + flush=True, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Sensitivity of the control +# -------------------------------------------------------------------------------------------------------------------- + + +def linearisation(env): + """Per-step (A, B) of the noise-free dynamics at the target (Euler step of the vector field) and the target.""" + dt = env.unwrapped.dt + target = np.asarray(env.unwrapped.reference_point, dtype=np.float64) + d = len(target) + + def f(x, u): + return np.asarray(env.unwrapped.continuous_f(np.array([u]))(0, x), dtype=np.float64).reshape(-1) + + h = 1e-5 + J = np.stack([(f(target + h * e, 0.0) - f(target - h * e, 0.0)) / (2 * h) for e in np.eye(d)], axis=1) + b = (f(target, h) - f(target, -h)) / (2 * h) + return np.eye(d) + dt * J, dt * b.reshape(-1, 1), target + + +def lqr_controller(env, action_cost_scale): + """LQR controller of the linearisation at the target, u = clip(-K (x - x_ref)), and its gain K. + + The reference is the per-step (undiscounted) LQR with the environment's Q and the scaled R, a benchmark of what a + linear controller achieves near the target, not the optimum of the discounted nonlinear problem. + """ + A, B, target = linearisation(env) + Q = np.asarray(env.unwrapped.Q, dtype=np.float64) + R = np.asarray(env.unwrapped.R, dtype=np.float64) * action_cost_scale + P = solve_discrete_are(A, B, Q, R) + K = np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A) + low, high = env.action_space.low, env.action_space.high + return (lambda x: np.clip(-K @ (x - target), low, high)), K + + +def evaluate(env, controller, seed, episodes, horizon, near_radius): + """Benchmark return, its state and action parts, and how close the state stays to the target, per episode mean. + + Episode e is seeded with seed * 1000 + e, so every controller meets the same initial states (and, for the double + well, the same noise). The state and action costs use the benchmark weights Q and R. "Near" means a Euclidean + distance from the target below `near_radius`; the late distance is the mean distance over the last quarter. + """ + target = np.asarray(env.unwrapped.reference_point, dtype=np.float64) + Q, R = np.asarray(env.unwrapped.Q), np.asarray(env.unwrapped.R) + totals = {"return": [], "state_cost": [], "action_cost": [], "late_distance": [], "fraction_near": []} + for episode in range(episodes): + np.random.seed(seed * 1000 + episode) + observation = env.reset() + total, state_cost, action_cost, distance = 0.0, 0.0, 0.0, [] + for _ in range(horizon): + x = np.asarray(observation, dtype=np.float64) + u = controller(x) + state_cost += float((x - target) @ Q @ (x - target)) + action_cost += float(u @ R @ u) + observation, reward, done, _ = env.step(u) + total += float(reward) + distance.append(float(np.linalg.norm(x - target))) + if done: + break + distance = np.array(distance) + totals["return"].append(total) + totals["state_cost"].append(state_cost) + totals["action_cost"].append(action_cost) + totals["late_distance"].append(float(distance[-len(distance) // 4 :].mean())) + totals["fraction_near"].append(float((distance < near_radius).mean())) + return {key: float(np.mean(values)) for key, values in totals.items()} + + +def gain_near_target(policy, env, samples=3000): + """Least-squares gain K of the mean action, u = -K (x - x_ref), on states near the target. + + The states are drawn uniformly from the cube of half-width 0.3 around the target with NumPy's global generator, + which `sensitivity_run` leaves in the state reached after the mean-action evaluation. This order is kept so that + the gains equal those of the scripts these checks were first run with. The drawn states, and hence the gain, + therefore depend on that call order and on the episode count. + """ + target = np.asarray(env.unwrapped.reference_point, dtype=np.float64) + act = deployed_controller(policy, mean=True) + states = target + np.random.uniform(-0.3, 0.3, size=(samples, len(target))) + actions = np.array([act(x) for x in states]) + return (-np.linalg.lstsq(states - target, actions, rcond=None)[0]).reshape(-1).tolist() + + +def sensitivity_run(env_id, seed, args): + """Sensitivity study for one environment and seed: the references, then every variant of VARIANTS[env_id].""" + start = time.time() + env = gym.make(env_id) + near = NEAR_RADIUS[env_id] + horizon = HORIZON[env_id] + result = { + "env": env_id, + "seed": seed, + "settings": { + "horizon": horizon, + "evaluation_episodes": args.evaluation_episodes, + "near_radius": near, + "variants": VARIANTS[env_id], + "variant_subset": args.variants, + "lqr_references": LQR_REFERENCES[env_id], + }, + "references": {}, + "variants": {}, + } + + references = [("no control", lambda x: np.zeros(1), None)] + for name, scale in LQR_REFERENCES[env_id]: + controller, K = lqr_controller(env, scale) + references.append((name, controller, K)) + for name, controller, K in references: + summary = evaluate(env, controller, seed, args.evaluation_episodes, horizon, near) + result["references"][name] = {"evaluation": summary, "gain": None if K is None else K.ravel().tolist()} + + for name, order, scale, rounds, used_for in VARIANTS[env_id]: + if args.variants and name not in args.variants: + continue + env, tensor, policy, config = train_skvi(env_id, seed, state_order=order, action_cost_scale=scale) + transitions = None + for round_ in range(rounds): # re-identify on the random-agent data plus all transitions of the policies so far + paths, steps = max(1, config["num-paths"] // 2), config["num-steps-per-path"] + new = policy_transitions(env, policy, paths, steps, seed + 50_000 + round_) + transitions = ( + new if transitions is None else tuple(np.concatenate([a, b]) for a, b in zip(transitions, new)) + ) + env, tensor, policy, config = train_skvi( + env_id, + seed, + state_order=order, + action_cost_scale=scale, + extra_transitions=transitions, + extra_transitions_for=used_for, + ) + _, powers = monomials(env.observation_space.shape[0], order) + dynamics = local_dynamics(env, tensor, linear_indices(powers)) + entry = { + "state_order": order, + "action_cost_scale": scale, + "reidentification_rounds": rounds, + "policy_transitions_used_for": used_for, + "policy_transitions": 0 if transitions is None else int(len(transitions[0])), + "sampled": evaluate(env, deployed_controller(policy), seed, args.evaluation_episodes, horizon, near), + "mean_action": evaluate( + env, deployed_controller(policy, mean=True), seed, args.evaluation_episodes, horizon, near + ), + "gain": gain_near_target(policy, env), + "A_model": dynamics["A_model"], + } + result["variants"][name] = entry + print( + f"{env_id} seed {seed} {name:>18}: return {entry['sampled']['return']:.4g} (mean action " + f"{entry['mean_action']['return']:.4g}), near target {100 * entry['sampled']['fraction_near']:.0f}%, " + f"gain {np.round(entry['gain'], 2).tolist()}", + flush=True, + ) + result["seconds"] = time.time() - start + write(args.output_dir, "sensitivity", env_id, seed, result) + + +# -------------------------------------------------------------------------------------------------------------------- +# Results and summaries +# -------------------------------------------------------------------------------------------------------------------- + + +def write(output_dir, study, env_id, seed, result): + """Write one run, refusing to replace a file of the same run made with other settings (e.g. another variant set).""" + path = os.path.join(output_dir, f"{study}_{env_id.split('-')[0]}_{seed}.json") + if os.path.exists(path): + with open(path) as f: + previous = json.load(f).get("settings") + if json.dumps(previous, sort_keys=True) != json.dumps(result.get("settings"), sort_keys=True): + raise FileExistsError(f"{path} holds a run made with other settings; use another --output_dir") + with open(path, "w") as f: + json.dump(result, f, indent=1) + + +def check_variant_names(names, environments): + """Raise if a name passed to --variants is not a variant of the sensitivity study for these environments.""" + known = {name for environment in environments for name, *_ in VARIANTS.get(f"{environment}-v0", [])} + unknown = sorted(set(names) - known) + if unknown: + raise ValueError(f"unknown variants {unknown}; known: {sorted(known)}") + + +def load_results(output_dir, study, environment): + """All runs of one study and environment in `output_dir`, which must share their settings.""" + runs = [json.load(open(p)) for p in sorted(glob.glob(os.path.join(output_dir, f"{study}_{environment}_*.json")))] + settings = {json.dumps(run.get("settings"), sort_keys=True) for run in runs} + if len(settings) > 1: + raise ValueError(f"{output_dir} holds {study} runs of {environment} made with different settings") + return runs + + +def box_corner_distance(env_id): + """Distance from the target to the farthest corner of the environment's state box.""" + env = gym.make(env_id) + target = np.asarray(env.unwrapped.reference_point, dtype=np.float64) + low, high = env.observation_space.low, env.observation_space.high + return float(np.linalg.norm(np.maximum(np.abs(low - target), np.abs(high - target)))) + + +def summarize(output_dir): + """Print the medians over seeds of both studies, with the failure counts and the model's local coefficient.""" + for environment in ACCURACY_ENVIRONMENTS: + runs = load_results(output_dir, "accuracy", environment) + if not runs: + continue + + def median(*keys): + values = [] + for run in runs: + value = run + for key in keys: + value = value[key] + values.append(value) + return float(np.median(values)) + + print( + f"Accuracy, {environment} ({len(runs)} seed{'s' if len(runs) != 1 else ''}): model error " + f"{median('on_policy', 'common_scale', 'model'):.2g} along the policy, " + f"{median('random_agent', 'common_scale', 'model'):.2g} on random-agent data; persistence error " + f"{median('on_policy', 'common_scale', 'persistence'):.2g} and " + f"{median('random_agent', 'common_scale', 'persistence'):.2g}; state error along the policy " + f"{median('on_policy', 'state', 'error_rms'):.2g} against a change of " + f"{median('on_policy', 'state', 'change_rms'):.2g}; local Jacobian: largest entry error " + f"{max(r['local_dynamics']['max_abs_entry_error'] for r in runs):.1e}, relative error " + f"{max(r['local_dynamics']['relative_error'] for r in runs):.1e} (worst seed)" + ) + + for environment in SENSITIVITY_ENVIRONMENTS: + runs = load_results(output_dir, "sensitivity", environment) + if not runs: + continue + print(f"Sensitivity, {environment} ({len(runs)} seed{'s' if len(runs) != 1 else ''}): medians over seeds") + scale = 0.01 # R * dt (dt = 0.01): the double-well costs are also reported with actions this much cheaper + rows = [(name, [r["references"][name]["evaluation"] for r in runs]) for name in runs[0]["references"]] + for name in runs[0]["variants"]: + rows.append((f"SKVI {name}, deployed", [r["variants"][name]["sampled"] for r in runs])) + rows.append((f"SKVI {name}, mean action", [r["variants"][name]["mean_action"] for r in runs])) + for name, evaluations in rows: + cost = np.median([-e["return"] for e in evaluations]) + cheap = np.median([e["state_cost"] + scale * e["action_cost"] for e in evaluations]) + state = np.median([e["state_cost"] for e in evaluations]) + near = [100 * e["fraction_near"] for e in evaluations] + over_half = sum(v > 50 for v in near) + extra = f", cost under R*dt {cheap:,.0f}" if environment == "DoubleWell" else "" + print( + f" {name:<34} cost {cost:>10,.0f}{extra}, state cost {state:,.0f}, near target {np.median(near):.0f}%" + f" (seeds over half the episode: {over_half} of {len(evaluations)})" + ) + corner = box_corner_distance(f"{environment}-v0") + for name in runs[0]["variants"]: + mean_action = [r["variants"][name]["mean_action"] for r in runs] + failed = sum(e["fraction_near"] < FAILURE_FRACTION for e in mean_action) + left_box = sum(e["late_distance"] > corner for e in mean_action) + print( + f" {name:<34} mean action: seeds near the target under {100 * FAILURE_FRACTION:g}% of the episode " + f"{failed} of {len(runs)}, seeds whose late " + f"distance exceeds the farthest corner of the state box ({corner:.2f}) {left_box} of {len(runs)}" + ) + for name in runs[0]["variants"]: + gains = np.array([r["variants"][name]["gain"] for r in runs]) + first = np.array([r["variants"][name]["A_model"][0][0] for r in runs]) + print( + f" {name:<34} gain median {np.round(np.median(gains, axis=0), 2).tolist()}, " + f"model coefficient A[0, 0] {first.min():.4f} to {first.max():.4f}" + ) + + +if __name__ == "__main__": + args = ArgumentParser().parse_args() + if args.study not in ("accuracy", "sensitivity"): + raise ValueError("--study must be 'accuracy' or 'sensitivity'") + os.makedirs(args.output_dir, exist_ok=True) + torch.set_num_threads(4) + if not args.summarize_only: + defaults = ACCURACY_ENVIRONMENTS if args.study == "accuracy" else SENSITIVITY_ENVIRONMENTS + run = accuracy_run if args.study == "accuracy" else sensitivity_run + if args.variants: + check_variant_names(args.variants, args.environments or defaults) + for environment in args.environments or defaults: + for seed in args.seeds: + run(f"{environment}-v0", seed, args) + summarize(args.output_dir) diff --git a/pyproject.toml b/pyproject.toml index 38718b1..a3596fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,7 @@ ignore = ["F401", "E731"] [tool.pytest.ini_options] minversion = "8.0" +pythonpath = ["."] testpaths = ["tests"] python_files = ["test_*.py"] python_classes = ["Test*"] diff --git a/tests/test_koopman_prediction_validation.py b/tests/test_koopman_prediction_validation.py new file mode 100644 index 0000000..3b3af6a --- /dev/null +++ b/tests/test_koopman_prediction_validation.py @@ -0,0 +1,82 @@ +"""koopmanrl_utils.koopman_prediction_validation uses the tuned settings and measures what it says it measures.""" + +import contextlib +import io +import json +import os + +import numpy as np + +import koopmanrl_utils.koopman_prediction_validation as validation + +CONFIG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "configurations") +CONFIG_FILES = { + "LinearSystem-v0": "skvi_linear_system_hparams.json", + "Lorenz-v0": "skvi_lorenz_hparams.json", + "FluidFlow-v0": "skvi_fluid_flow_hparams.json", + "DoubleWell-v0": "skvi_double_well_hparams.json", +} + + +def small_run(monkeypatch, env_id, **overrides): + """One seed of `run_env` with a short horizon, few rollouts and one subsample size.""" + monkeypatch.setattr(validation, "HORIZON", 5) + monkeypatch.setattr(validation, "N_ROLLOUT_PATHS", 4) + monkeypatch.setattr(validation, "N_MC", 8) + monkeypatch.setattr(validation, "N_ID_GRID", [250]) + config = {**validation.CONFIGS[env_id], "num_paths": 10, "steps": 50, **overrides} + with contextlib.redirect_stdout(io.StringIO()): + return validation.run_env(env_id, config, seed=0) + + +def test_settings_match_the_tuned_configurations(): + for env_id, settings in validation.CONFIGS.items(): + with open(os.path.join(CONFIG_DIR, CONFIG_FILES[env_id])) as f: + tuned = json.load(f) + assert settings["num_paths"] == tuned["num-paths"] + assert settings["steps"] == tuned["num-steps-per-path"] + assert settings["so"] == tuned["state-order"] + assert settings["ao"] == tuned["action-order"] + + +def test_nrmse_is_relative_to_the_reference(): + reference = np.array([[3.0, 0.0], [4.0, 5.0]]) + assert validation.nrmse(reference, reference) == 0.0 + assert np.isclose(validation.nrmse(np.zeros_like(reference), reference), 1.0) + assert np.isclose(validation.nrmse(1.1 * reference, reference), 0.1) + + +def test_lifted_metrics_scale_each_function_and_drop_constants(): + rng = np.random.RandomState(0) + reference = np.vstack([np.ones(200), rng.randn(200), 1e3 * rng.randn(200)]) + base = reference + np.vstack([np.zeros(200), 0.2 * reference[1], 0.2 * reference[2]]) + prediction = reference + np.vstack([5.0 * np.ones(200), 0.1 * reference[1], 0.1 * reference[2]]) + metrics = validation.lifted_metrics(prediction, reference, base) + # the error of the constant function is ignored, and the large function does not dominate + assert np.isclose(metrics["err"], 0.1) + assert np.isclose(metrics["persistence"], 0.2) + assert np.isclose(metrics["increment"], 0.5) + assert metrics["noise_floor"] == 0.0 + + +def test_linear_system_is_predicted_to_rounding(monkeypatch): + onestep, multistep, sweep, _ = small_run(monkeypatch, "LinearSystem-v0") + assert onestep["phi_heldout"] < 1e-10 + assert onestep["heldout"] < 1e-10 + assert onestep["phi_persistence"] > 1e-2 + assert multistep.shape == (5, 3) and np.all(multistep[:, :2] < 1e-8) + assert sweep.shape == (1, 2) + + +def test_first_rollout_step_is_the_same_for_both_rollouts(monkeypatch): + # iterating K^u in dictionary space and re-evaluating the dictionary differ only from the second step on + _, multistep, _, dt = small_run(monkeypatch, "Lorenz-v0") + assert np.isclose(multistep[0, 0], multistep[0, 1], rtol=1e-9) + assert dt == 0.01 + + +def test_double_well_reports_a_noise_floor(monkeypatch): + # also runs the check of the vectorised drift against the environment's + onestep, multistep, _, _ = small_run(monkeypatch, "DoubleWell-v0") + assert onestep["phi_noise_floor"] > 0.0 and onestep["noise_floor"] > 0.0 + assert np.all(np.isfinite(multistep)) diff --git a/tests/test_skvi_policy_checks.py b/tests/test_skvi_policy_checks.py new file mode 100644 index 0000000..f31b107 --- /dev/null +++ b/tests/test_skvi_policy_checks.py @@ -0,0 +1,61 @@ +"""The helpers of koopmanrl_utils.skvi_policy_checks reproduce the package's SKVI quantities.""" + +import contextlib +import io + +import gym +import numpy as np +import torch + +import koopmanrl.environments # noqa: F401 +from koopmanrl.soft_koopman_value_iteration import ( + DiscreteKoopmanValueIterationPolicy, + generate_koopman_tensor, +) +from koopmanrl_utils.skvi_policy_checks import ( + gibbs_policy, + monomials, + quadratic_cost, + quadratic_form, +) + + +def small_policy(env_id="LinearSystem-v0", seed=1): + np.random.seed(seed) + torch.manual_seed(seed) + with contextlib.redirect_stdout(io.StringIO()): + env = gym.make(env_id) + tensor = generate_koopman_tensor(env_id, seed, 5, 50, 2, 2, "ols") + actions = torch.from_numpy(np.linspace(env.action_space.low, env.action_space.high, 11)).T + policy = DiscreteKoopmanValueIterationPolicy( + env_id, 0.99, 1.0, tensor, actions, env.vectorized_cost_fn, seed, use_ols=True, dt=None + ) + policy.value_function_weights = torch.from_numpy(np.random.randn(tensor.phi_dim, 1)) + return env, policy + + +def test_quadratic_cost_matches_environment(): + env, policy = small_policy() + states = np.random.randn(4, 3) + expected = env.vectorized_cost_fn(torch.from_numpy(states), policy.all_actions.T).numpy() + actual = quadratic_cost(env)(states, policy.all_actions.T.numpy()).numpy() + assert np.allclose(actual, expected) + + +def test_gibbs_policy_matches_package(): + _, policy = small_policy() + states = np.array([[1.0, 2.0, -3.0], [10.0, -5.0, 0.5]]) + batched = gibbs_policy(policy, states) + for n, state in enumerate(states): + # the package's pis takes one state as a (1, d) row and returns identical columns + single = policy.pis(torch.from_numpy(state.reshape(1, -1))).detach().numpy()[:, 0] + assert np.allclose(batched[:, n], single, atol=1e-12) + + +def test_quadratic_form_reads_the_degree_two_terms(): + names, powers = monomials(2, 2) + coefficients = np.zeros(len(names)) + coefficients[names.index("x1^2")] = 3.0 + coefficients[names.index("x1x2")] = 4.0 + coefficients[names.index("x2^2")] = 5.0 + assert np.allclose(quadratic_form(coefficients, powers), [[3.0, 2.0], [2.0, 5.0]]) diff --git a/tests/test_skvi_sensitivity_checks.py b/tests/test_skvi_sensitivity_checks.py new file mode 100644 index 0000000..086d72e --- /dev/null +++ b/tests/test_skvi_sensitivity_checks.py @@ -0,0 +1,200 @@ +"""The helpers of koopmanrl_utils.skvi_sensitivity_checks reproduce the quantities they stand for.""" + +import contextlib +import io + +import gym +import numpy as np +import torch + +import koopmanrl.environments # noqa: F401 +from koopmanrl.soft_koopman_value_iteration import ( + DiscreteKoopmanValueIterationPolicy, + generate_koopman_tensor, +) +from koopmanrl_utils.skvi_policy_checks import ( + identify_tensor, + monomials, + quadratic_cost, +) +from koopmanrl_utils.skvi_sensitivity_checks import ( + check_variant_names, + double_well_conditional_mean, + evaluate, + features, + gibbs_sampler, + linear_indices, + local_dynamics, + lqr_controller, + one_step_errors, + predict, + write, +) + + +def small_tensor(env_id="LinearSystem-v0", seed=1, state_order=2, action_order=2): + np.random.seed(seed) + torch.manual_seed(seed) + with contextlib.redirect_stdout(io.StringIO()): + env = gym.make(env_id) + tensor = generate_koopman_tensor(env_id, seed, 5, 50, state_order, action_order, "ols") + return env, tensor + + +def small_policy(env_id="LinearSystem-v0", seed=1): + env, tensor = small_tensor(env_id, seed) + actions = torch.from_numpy(np.linspace(env.action_space.low, env.action_space.high, 11)).T + policy = DiscreteKoopmanValueIterationPolicy( + env_id, 0.99, 1.0, tensor, actions, quadratic_cost(env), seed, use_ols=True, dt=None + ) + policy.value_function_weights = torch.from_numpy(np.random.randn(tensor.phi_dim, 1)) + return env, policy + + +def test_gibbs_sampler_matches_package(): + _, policy = small_policy() + probabilities, _ = gibbs_sampler(policy) + for state in np.array([[1.0, 2.0, -3.0], [10.0, -5.0, 0.5]]): + single = policy.pis(torch.from_numpy(state.reshape(1, -1))).detach().numpy()[:, 0] + assert np.allclose(probabilities(state), single, atol=1e-12) + + +def test_predict_matches_tensor(): + _, tensor = small_tensor() + states = np.random.default_rng(0).normal(size=(4, 3)) + actions = np.array([-1.0, 0.0, 0.5, 2.0]) + Phi = features(tensor, states) + K_u = tensor.K_(torch.from_numpy(actions.reshape(1, -1)).to(tensor.K.dtype)).numpy().astype(np.float64) + assert np.allclose(predict(tensor, Phi, actions), np.einsum("nij,nj->ni", K_u, Phi), rtol=1e-5, atol=1e-8) + + +def test_double_well_conditional_mean_through_the_environment(): + # For a dictionary of degree two the Gaussian expectation equals the average over the four noise draws + # (+-1, +-1), which match the first and second moments; the draws go through the environment's own step. + env = gym.make("DoubleWell-v0") + env.reset() + _, powers = monomials(2, 2) + X = np.array([[0.3, -0.4], [-1.2, 0.8]]) + U = np.array([1.5, -2.0]) + exact, mean_state = double_well_conditional_mean(env, X, U, powers) + for n, (x, u) in enumerate(zip(X, U)): + next_states = [] + for draw in ([1.0, 1.0], [1.0, -1.0], [-1.0, 1.0], [-1.0, -1.0]): + env.unwrapped.step_count = 0 + env.unwrapped.random_draws[0] = np.array(draw).reshape(2, 1) + next_states.append(env.unwrapped.f(x, np.array([u]))) + next_states = np.array(next_states) + average = np.prod(next_states[:, :, None] ** powers[None, :, :], axis=1).mean(axis=0) + assert np.allclose(exact[n], average, atol=1e-12) + assert np.allclose(mean_state[n], next_states.mean(axis=0), atol=1e-12) + + +def test_one_step_errors_at_a_common_scale(): + rng = np.random.default_rng(0) + target = rng.normal(size=(50, 3)) * np.array([1.0, 1.0, 10.0]) + target[:, 0] = 1.0 # constant dictionary function, excluded from the error + prediction = target + 0.1 + base = target + 1.0 + X, next_state = rng.normal(size=(50, 2)), rng.normal(size=(50, 2)) + nonconstant = np.array([False, True, True]) + out = one_step_errors(prediction, target, base, X, next_state, nonconstant, [1, 2], common_scale=np.ones(2)) + assert np.isclose(out["common_scale"]["model"], 0.1) + assert np.isclose(out["common_scale"]["persistence"], 1.0) + assert np.isclose(out["common_scale"]["ratio"], 0.1) + own = np.sqrt(np.mean(target[:, 1:] ** 2, axis=0)) + assert np.isclose(out["own_scale"]["model"], np.sqrt(np.mean((0.1 / own) ** 2))) + # the state error compares the predicted degree-one functions (here columns 1 and 2) with the next state + error = np.linalg.norm(prediction[:, [1, 2]] - next_state, axis=1) + assert np.isclose(out["state"]["error_rms"], np.sqrt(np.mean(error**2))) + assert np.isclose(out["state"]["change_rms"], np.sqrt(np.mean(np.linalg.norm(next_state - X, axis=1) ** 2))) + + +def test_local_dynamics_of_the_linear_system(): + env, tensor = small_tensor() + _, powers = monomials(3, 2) + dynamics = local_dynamics(env, tensor, linear_indices(powers)) + assert np.allclose(dynamics["A_true"], env.unwrapped.A, atol=1e-8) + assert dynamics["max_abs_entry_error"] < 1e-3 # the order-2 dictionary contains the linear dynamics exactly + + +def test_lqr_controller_stabilises_the_linearisation(): + env = gym.make("DoubleWell-v0") + for scale in (1.0, 0.01): + _, K = lqr_controller(env, scale) + dt = env.unwrapped.dt + A = np.eye(2) + dt * np.array([[4.0, 0.0], [0.0, -2.0]]) # drift 4x - 4x^3 and -2y at the origin + B = dt * np.ones((2, 1)) + assert np.abs(np.linalg.eigvals(A - B @ K)).max() < 1 + + +def test_action_cost_scale(): + env = gym.make("DoubleWell-v0") + states, actions = np.random.default_rng(0).normal(size=(5, 2)), np.linspace(-3, 3, 7).reshape(-1, 1) + assert torch.equal(quadratic_cost(env, 1.0)(states, actions), quadratic_cost(env)(states, actions)) + cheap = quadratic_cost(env, 0.01)(states, actions) - quadratic_cost(env)(states, actions) + assert np.allclose(cheap.numpy(), -0.99 * np.asarray(env.unwrapped.R).item() * actions**2) + + +class ScalarSystem: + """Deterministic test system x' = x / 2 + u with cost x^2 + 2 u^2, starting at x = 1.""" + + def __init__(self): + self.unwrapped = self + self.Q, self.R, self.reference_point = np.eye(1), 2 * np.eye(1), np.zeros(1) + + def reset(self): + self.x = np.ones(1) + return self.x.copy() + + def step(self, u): + reward = -float(self.x @ self.Q @ self.x + u @ self.R @ u) + self.x = self.x / 2 + u + return self.x.copy(), reward, False, {} + + +def test_evaluate_sums_costs_and_measures_the_distance(): + env = ScalarSystem() + free = evaluate(env, lambda x: np.zeros(1), seed=0, episodes=2, horizon=4, near_radius=0.3) + # states 1, 1/2, 1/4, 1/8: state cost 1 + 1/4 + 1/16 + 1/64, two of four steps within 0.3, last quarter 1/8 + assert np.isclose(free["state_cost"], 1.328125) and free["action_cost"] == 0 + assert np.isclose(free["return"], -1.328125) + assert free["fraction_near"] == 0.5 and np.isclose(free["late_distance"], 0.125) + # u = -x/2 sends the state to the target in one step: state cost 1, action cost 2 (1/2)^2 + controlled = evaluate(env, lambda x: -x / 2, seed=0, episodes=1, horizon=4, near_radius=0.3) + assert np.isclose(controlled["state_cost"], 1.0) and np.isclose(controlled["action_cost"], 0.5) + assert np.isclose(controlled["return"], -1.5) and controlled["fraction_near"] == 0.75 + + +def test_identify_tensor_appends_transitions_in_single_precision(): + config = {"num-paths": 3, "num-steps-per-path": 20, "state-order": 2, "action-order": 2} + rng = np.random.default_rng(0) + extra = (rng.normal(size=(5, 2)), rng.normal(size=(5, 1)), rng.normal(size=(5, 2))) + with contextlib.redirect_stdout(io.StringIO()): + base = identify_tensor("DoubleWell-v0", 1, config, 0.01) + refitted = identify_tensor("DoubleWell-v0", 1, config, 0.01, extra) + for name, appended in zip(("X", "U", "Y"), extra): + stored = getattr(refitted, name).numpy() + assert stored.shape[1] == getattr(base, name).shape[1] + 5 + assert np.array_equal(stored[:, : getattr(base, name).shape[1]], getattr(base, name).numpy()) + assert np.array_equal(stored[:, -5:], appended.T.astype(np.float32).astype(stored.dtype)) + + +def test_runs_with_other_settings_are_not_overwritten(tmp_path): + write(str(tmp_path), "sensitivity", "Lorenz-v0", 100, {"settings": {"variant_subset": []}}) + write(str(tmp_path), "sensitivity", "Lorenz-v0", 100, {"settings": {"variant_subset": []}}) # a rerun is fine + try: + write(str(tmp_path), "sensitivity", "Lorenz-v0", 100, {"settings": {"variant_subset": ["order 3"]}}) + except FileExistsError: + pass + else: + raise AssertionError("a run with other settings replaced the file") + + +def test_unknown_variant_names_are_rejected(): + check_variant_names(["order 3, refit x1, tensor only"], ["Lorenz"]) + try: + check_variant_names(["order 3, refit"], ["Lorenz"]) + except ValueError: + pass + else: + raise AssertionError("an unknown variant name was accepted")