#!/usr/bin/env python # -*- coding: utf-8 -*- # File: physics.py # # Physics module for: # "Physics-Informed Communicative Reinforcement Learning via Hamilton-Jacobi # Optimal Control for Robust Anatomical Landmark Detection in Low-Res Brain MRI" # # WHAT THIS IS (and what the two source PDFs got wrong): # ----------------------------------------------------------------------------- # The proposal calls the constraint ||grad V|| * F = 1 a "Hamilton-Jacobi-Bellman" # equation and claims it is "the same principle that underlies RL". That is not # quite right. This equation is the *stationary Hamilton-Jacobi (Eikonal)* PDE. It # is the HJ equation for the minimum-arrival-time / geodesic problem, NOT the # Hamilton-Jacobi-*Bellman* equation whose solution is the RL value function. # They are cousins, not twins (see EVALUATION.md). # # The *correct and defensible* way to use it here is: # 1. Precompute, ONCE per image (offline), the geodesic navigation potential V # by solving the Eikonal equation from each landmark, with an edge-modulated # speed field F. (Solving a PINN per RL step, as the PDF suggests, is both # unnecessary and prohibitively expensive.) # 2. Use -V as a *potential-based reward shaping* term (Ng, Harada & Russell, # 1999). This provably leaves the optimal policy unchanged while giving the # agent a dense, anatomically-aware learning signal. The PDF's ad-hoc # R = -dd - lambda||grad V|| + gamma C is NOT policy-invariant and can bias # the optimum. # 3. Optionally expose V (and other derived maps) as extra input channels. # # This module implements 1-3 with pure numpy/scipy/skimage and is fully runnable # and cache-backed. No autograd PINN is required. # ----------------------------------------------------------------------------- import os import hashlib import numpy as np from scipy import ndimage as ndi try: from skimage.graph import MCP_Geometric _HAVE_SKIMAGE = True except Exception: # pragma: no cover _HAVE_SKIMAGE = False # ============================================================================= # Low-level field operators # ============================================================================= def _normalize(vol, lo_pct=1.0, hi_pct=99.0): vol = vol.astype(np.float32) lo = np.percentile(vol, lo_pct) hi = np.percentile(vol, hi_pct) return np.clip((vol - lo) / (hi - lo + 1e-6), 0.0, 1.0) def edge_strength(vol, sigma=1.0): """G(x) = |grad I| of a lightly smoothed, normalized volume.""" im = ndi.gaussian_filter(_normalize(vol), sigma) grads = np.gradient(im) return np.sqrt(sum(g ** 2 for g in grads)).astype(np.float32) def speed_field(vol, beta=8.0, sigma=1.0): """F(x) = 1 / (1 + beta |grad I|), in (0, 1]. Fast (F->1) in smooth tissue, slow (F->0) at anatomical boundaries. This is the standard geodesic-active-contour speed and is bounded, unlike the PDF's alternative F = exp(-alpha G) which is also fine but less numerically tame. """ g = edge_strength(vol, sigma) g = g / (g.max() + 1e-6) return (1.0 / (1.0 + beta * g)).astype(np.float32) def hessian_vesselness(vol, sigma=1.5): """Cheap Frangi-like ridge/vesselness proxy from Hessian eigenvalues. NOTE: for brain landmark detection this channel is of dubious value (it was designed for tubular vessels); included for completeness / ablation only. """ from skimage.feature import hessian_matrix, hessian_matrix_eigvals n = _normalize(vol) H = hessian_matrix(n, sigma=sigma, use_gaussian_derivatives=False) eigs = hessian_matrix_eigvals(H) l_small = np.abs(eigs[-1]) l_large = np.abs(eigs[0]) + 1e-6 return np.clip(l_small / l_large, 0, 3).astype(np.float32) # ============================================================================= # The Eikonal / stationary Hamilton-Jacobi solver # ============================================================================= def geodesic_potential(vol, seed_xyz, beta=8.0, sigma=1.0): """Solve ||grad V|| * F = 1 from `seed_xyz` (a landmark). Implemented as a grid minimum-cost distance with slowness cost = 1/F, which is exactly the first-order discretization of the Eikonal equation (this is what fast-marching computes). Returns the geodesic arrival-time field V. Falls back to a plain distance transform if scikit-image is unavailable. """ seed = tuple(int(round(c)) for c in seed_xyz) if not _HAVE_SKIMAGE: m = np.ones(vol.shape, dtype=bool) m[seed] = False return ndi.distance_transform_edt(m).astype(np.float32) F = speed_field(vol, beta, sigma) cost = (1.0 / (F + 1e-6)).astype(np.float64) mcp = MCP_Geometric(cost) V, _ = mcp.find_costs([seed]) V = np.asarray(V, dtype=np.float32) finite = np.isfinite(V) V[~finite] = V[finite].max() if finite.any() else 0.0 return V def euclidean_potential(shape, seed_xyz, spacing=(1, 1, 1)): """d(x): straight-line distance to the landmark. This is what the PDF's L_Consistency = ||V - d||^2 pulls V toward. Because V is *geodesic*, d != V wherever F varies, so that loss fights the Eikonal constraint. Use `geodesic` targets for consistency instead (see below). """ m = np.ones(shape, dtype=bool) m[tuple(int(round(c)) for c in seed_xyz)] = False return ndi.distance_transform_edt(m, sampling=spacing).astype(np.float32) # ============================================================================= # Per-image physics provider (cached) -- the object the environment uses # ============================================================================= class PhysicsProvider: """Precomputes and caches per-(image, landmark) physics fields. Designed to be created once per environment and queried each step. All heavy computation happens on the first touch of a given (image, landmark) pair. """ def __init__(self, beta=8.0, sigma=1.0, cache_dir=None, channels=("mri", "grad", "potential", "dist")): self.beta = beta self.sigma = sigma self.cache_dir = cache_dir self.channels = tuple(channels) self._mem = {} # in-process cache if cache_dir: os.makedirs(cache_dir, exist_ok=True) # ---- keys / disk cache ------------------------------------------------ def _key(self, image_id, seed_xyz): h = hashlib.md5(f"{image_id}|{seed_xyz}|{self.beta}|{self.sigma}".encode()).hexdigest()[:16] return h def _load_disk(self, key): if not self.cache_dir: return None p = os.path.join(self.cache_dir, key + ".npz") if os.path.exists(p): z = np.load(p) return {k: z[k] for k in z.files} return None def _save_disk(self, key, fields): if not self.cache_dir: return np.savez_compressed(os.path.join(self.cache_dir, key + ".npz"), **fields) # ---- main API --------------------------------------------------------- def fields_for(self, volume, seed_xyz, image_id="img"): """Return dict with 'V' (geodesic potential), 'F' (speed), 'd' (euclid).""" key = self._key(image_id, tuple(map(int, seed_xyz))) if key in self._mem: return self._mem[key] disk = self._load_disk(key) if disk is not None: self._mem[key] = disk return disk F = speed_field(volume, self.beta, self.sigma) V = geodesic_potential(volume, seed_xyz, self.beta, self.sigma) d = euclidean_potential(volume.shape, seed_xyz) fields = {"V": V, "F": F, "d": d} self._save_disk(key, fields) self._mem[key] = fields return fields def channel_stack(self, volume, seed_xyz, image_id="img"): """Build the multi-channel physics tensor requested in Step 1 of the PDF. Returns an array of shape (C, X, Y, Z), normalized to [0,1] per channel. Redundant channels (e.g. grad vs potential) are correlated ~0.9; keep the set small for real training. Default: MRI, grad, potential, dist. """ f = self.fields_for(volume, seed_xyz, image_id) maps = { "mri": _normalize(volume), "grad": _norm01(edge_strength(volume, self.sigma)), "vesselness": _norm01(hessian_vesselness(volume)), "potential": _norm01(f["V"]), "dist": _norm01(f["d"]), "speed": f["F"], } return np.stack([maps[c] for c in self.channels], axis=0).astype(np.float32) # ---- reward shaping ---------------------------------------------------- def shaping_potential(self, volume, seed_xyz, loc, image_id="img"): """Phi(s) = -V(loc): the potential for Ng-1999 potential-based shaping. The shaped reward is r' = r + gamma*Phi(s') - Phi(s), which is guaranteed to preserve the optimal policy while densifying the signal. Use this INSTEAD of adding raw ||grad V|| to the reward. """ V = self.fields_for(volume, seed_xyz, image_id)["V"] i = tuple(np.clip(np.round(loc).astype(int), 0, np.array(V.shape) - 1)) return -float(V[i]) def _norm01(a): a = a.astype(np.float32) return (a - a.min()) / (a.max() - a.min() + 1e-6) # ============================================================================= # Self-test (runs on any 3D numpy volume) # ============================================================================= if __name__ == "__main__": import sys if len(sys.argv) > 1: vol = np.load(sys.argv[1]) else: # synthetic blob volume z, y, x = np.mgrid[0:60, 0:60, 0:60] vol = (np.sin(x / 5.0) * np.cos(y / 5.0) * 100 + 200).astype(np.float32) if vol.ndim == 4: vol = vol[..., 0] seed = tuple(np.array(vol.shape) // 2) pp = PhysicsProvider() f = pp.fields_for(vol, seed) print("V range", float(f["V"].min()), float(f["V"].max())) print("F range", float(f["F"].min()), float(f["F"].max())) stk = pp.channel_stack(vol, seed) print("channel stack", stk.shape) print("Phi at a random loc:", pp.shaping_potential(vol, seed, [10, 10, 10])) print("physics.py self-test OK")