diff --git a/.pi-lens.json b/.pi-lens.json new file mode 100644 index 0000000..7a7521b --- /dev/null +++ b/.pi-lens.json @@ -0,0 +1,3 @@ +{ + "format": { "enabled": false } +} diff --git a/AGENTS.md b/AGENTS.md index b69e092..7d7779a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,7 +54,7 @@ joint space**. - Demos: `python -m demos.demo_*`. - **Docs** (API reference is generated, not hand-written): - `make docs` — install the docs toolchain (`requirements-docs.txt`). - - `make docs-serve` — regenerate reference + serve at http://127.0.0.1:8000. + - `make docs-serve` — regenerate reference + serve at . - `make docs-build` — regenerate reference + build static `site/`. - `scripts/gen_api.py` walks each subpackage's `__all__` and emits one page per subpackage plus the nav — **adding an export is all it takes** for it to @@ -75,6 +75,11 @@ joint space**. path) then `scripts/build_scenario.py` (zig build + oracle + stamp + verify). The scenario TOML's `[compile]` section (`n_x`, `n_u`, `optimize`, `target`, `solver_dir`) is the build spec; never clobbers `src/shinro/runtime/graph_data.zig`. +- `make measure-kernels` — the kernel **size metric** (`shinro/codegen/measure.py` + plus a thin `scripts/` shim): the C-ABI host buffers, the VM's internal stack + buffer (`buf: [buf_len]f64`), and — with `BUILD=1` — the compiled artifact + bytes and compile cost. `DIMS=3x3x6x3,...` sweeps `D_x x D_u x N x K` rollout + shapes (multi-input systems included); static-only by default, so no compiler. ## Zig lowering (codegen → `.so`) @@ -96,7 +101,7 @@ joint space**. ## Key Files | File | Purpose | -|------|---------| +| ------ | --------- | | `src/shinro/components.py` | The five ABCs | | `src/shinro/codegen/` | Trace → compose → interpret → lower pipeline; `lower_zig.py` serializes a graph to `src/shinro/runtime/graph_data.zig` | | `src/shinro/factories/registry.py` | Component registry + config-driven factory | diff --git a/Makefile b/Makefile index b9db468..2d49545 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ test-controllers test-estimators test-plants test-trajectories test-armrobot \ test-components test-array-backend test-batched-adapter test-controllability test-factories \ test-linearization test-adversarial test-mcp-server test-mcp-functional \ - compile + compile measure-kernels # Install the package in editable mode install: @@ -125,6 +125,18 @@ compile: python3 scripts/gen_scenario.py $(SCENARIO) --out $(OUT) python3 scripts/build_scenario.py $(OUT) --scenario $(SCENARIO) $(FLAGS) +# ─────────────────────────────────────────────────────────────────────────── +# Measure lowered-kernel sizes as a metrics document: the C-ABI host buffers, +# the VM's internal stack buffer, the compiled artifact, and the compile cost. +# DIMS is a comma-separated list of `D_x x D_u x N x K` rollout shapes; +# BUILD=1 compiles (minutes at production sizes) and JSON= writes the +# machine-readable document. Static-only is fast (no compiler). +# ─────────────────────────────────────────────────────────────────────────── +DIMS ?= 3x3x6x3,3x3x100x15,6x6x10x4,8x4x12x5 +OPTIMIZE ?= ReleaseFast +measure-kernels: + python3 scripts/measure_kernels.py --dims $(DIMS) --optimize $(OPTIMIZE) $(if $(JSON),--json $(JSON),) $(if $(BUILD),--build,) + # Run linter and type checker lint: ruff check . diff --git a/README.md b/README.md index 31f6ac0..228a716 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,7 @@ Non-conforming commits land under **Other** (or are dropped). `chore: release vX.Y.Z` commits are skipped automatically. Semantic versioning: + - `patch` — backwards-compatible fix (`v0.1.0` → `v0.1.1`) - `minor` — backwards-compatible feature (`v0.1.1` → `v0.2.0`) - `major` — incompatible API change (`v0.2.0` → `v1.0.0`) @@ -112,6 +113,11 @@ python -m demos.demo_arm_trajectory # arm trajectory + liv python -m demos.demo_base_tracking # base tracking, LQR + observer python -m demos.demo_base_tracking --controller mpc # base tracking, MPC python -m demos.demo_pick_and_place # full pick-and-place sequence +python -m demos.demo_mppi # MPPI model wiring + lowering contract +python -m demos.demo_mppi_nonlinear # MPPI on a nonlinear plant, lowered +python -m demos.demo_mppi_nonlinear --build # ...and run the compiled kernel +python -m demos.demo_smc # SMC on a nonlinear plant (f/g host-side) +python -m demos.demo_smc --build # ...and run the compiled kernel ``` Auto-generate a robot config from a MuJoCo model: diff --git a/demos/demo_mppi.py b/demos/demo_mppi.py new file mode 100644 index 0000000..ec21c3f --- /dev/null +++ b/demos/demo_mppi.py @@ -0,0 +1,244 @@ +# FILE: demos/demo_mppi.py +"""MPPI demo: how to wire ``dynamics_fn`` / ``cost_fn``, and the lowering contract. + +Four ways exist to give MPPI its model, and which one you pick decides whether +the controller can be compiled to a Zig kernel: + +1. **Constructor injection** — ``MPPIController(dynamics_fn=..., cost_fn=...)`` + for a hand-written model. +2. **Attribute injection** — set ``ctrl.dynamics_fn`` / ``ctrl.cost_fn`` after + construction (the ``from_config`` path). +3. **Plant-driven** — ``ctrl.attach_plant(plant, Q=..., R=...)`` builds both + from the plant's model; this is what ``ScenarioFactory`` does. +4. **Config-driven** — ``MPPIController.from_config(cfg)`` (TOML cannot hold a + callable, so the callables are injected afterwards by one of the above). + +The callables are **batched** and **traced through** (they are inlined into the +lowered graph), so for a lowerable controller they must be trace-safe: operator +and ``bk``-op only, no raw numpy, no ``x[i]`` indexing, no ``bk.sum``, and no +rank-differing broadcast. ``attach_plant``'s LTI path already satisfies this. + +Everything here runs on the default install — no MuJoCo, no torch required. + +Usage: + python -m demos.demo_mppi +""" + +from __future__ import annotations + +import numpy as np + +from shinro.codegen import interpret +from shinro.codegen.trace_node import trace_node +from shinro.controllers.mppi import MPPIController +from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot +from shinro.utils.array_backend import NumpyBackend + +DT = 0.02 +N_SAMPLES = 200 +HORIZON = 15 + +# A first-order integrator, x' = x + dt*u, as diagonal weights. These are the +# model constants a user would normally derive from their own system. +D_X, D_U = 1, 1 +Q_DIAG = np.array([1.0]) # (D_x,) state weights +R_DIAG = np.array([0.1]) # (D_u,) control weights + + +# ─── the two callables, written trace-safely ─────────────────────────────── + + +def integrator_dynamics(x_batch, u_batch, dt): + """Batched dynamics: ``(N, D_x), (N, D_u)`` -> ``(N, D_x)``. + + Operator-only — no numpy, no ``x[i]`` indexing, no ``bk.sum`` — so this is + also what tracing records when the controller is lowered. + """ + return x_batch + dt * u_batch + + +def quadratic_cost(x_batch, u_batch): + """Batched stage cost: ``(N, D_x), (N, D_u)`` -> ``(N,)`` per-sample scalar. + + Each quadratic form is written as a contraction (``(N, D) @ (D,) -> (N,)``) + rather than ``bk.sum(x*x*W, axis=1)``: a sum over an axis is a matmul + identity, and this form traces (``TraceBackend`` has no ``sum`` op). + """ + return (x_batch * x_batch) @ Q_DIAG + (u_batch * u_batch) @ R_DIAG + + +def build_custom_controller(bk): + """Path 1 — constructor injection: a hand-written model.""" + return MPPIController( + dynamics_fn=integrator_dynamics, + cost_fn=quadratic_cost, + num_samples=N_SAMPLES, + temperature=1.0, + dt=DT, + horizon=HORIZON, + noise_sigma=[1.0], + seed=0, + backend=bk, + ) + + +# ─── 1. constructor injection ────────────────────────────────────────────── + + +def demo_custom_model(bk): + print("=== 1. Constructor injection: hand-written dynamics + cost ===") + ctrl = build_custom_controller(bk) + + x = np.array([1.0]) # (D_x,) + for step in range(200): + u = bk.to_numpy(ctrl.compute(x)) + x = x + DT * u # the plant the controller was told about + if step % 50 == 0: + print(f" step {step:3d}: x = {x[0]:+.4f} u = {u[0]:+.4f}") + print(f" final |x| = {abs(x[0]):.4f} (regulated toward the origin)\n") + + +# ─── 2. plant-driven (attach_plant) ──────────────────────────────────────── + + +def demo_plant_model(bk): + print("=== 2. Plant-driven: attach_plant(plant, Q=, R=) ===") + plant = HolonomicMobileRobot( + num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=DT, backend=bk + ) + ctrl = MPPIController( + num_samples=200, + temperature=1.0, + dt=DT, + horizon=10, + noise_sigma=[1.0, 1.0, 1.0], + seed=1, + backend=bk, + ) + # The plant's LTI model supplies dynamics (x @ A.T + u @ B.T); Q/R are the + # quadratic cost weights. attach_plant OVERWRITES an earlier cost_fn. + ctrl.attach_plant(plant, Q=np.array([10.0, 10.0, 10.0]), R=np.array([0.1, 0.1, 0.1])) + + x = np.zeros(3) + x_ref = np.array([1.0, 0.0, 0.0]) + for _ in range(300): + u = bk.to_numpy(ctrl.compute(bk.array(x), bk.array(x_ref))) # x_ref tracked + x = x + DT * u # A = I, B = dt*I for this plant + print(f" tracked x_ref = {x_ref}, reached x = {np.round(x, 4)}") + print(f" |x - x_ref| = {np.linalg.norm(x - x_ref):.4f}\n") + + +# ─── 3. from_config ──────────────────────────────────────────────────────── + + +def demo_config_model(bk): + print("=== 3. Config-driven: from_config() + injection ===") + config = { + "num_samples": 100, + "temperature": 1.0, + "dt": DT, + "horizon": 10, + "noise_sigma": [1.0], + "u_min": [-5.0], + "u_max": [5.0], + "seed": 7, + "state_cost": [1.0], # -> ctrl._Q + "control_cost": [0.1], # -> ctrl._R + } + ctrl = MPPIController.from_config(config, backend=bk) + # A callable cannot live in TOML: inject it (path 1 or 2), or call + # attach_plant(plant) to build it from the plant instead. + ctrl.dynamics_fn = integrator_dynamics + ctrl.cost_fn = quadratic_cost + u = bk.to_numpy(ctrl.compute(bk.array([1.0]))) + print(f" from_config + injection: u = {u[0]:+.4f} (bounded by u_min/u_max)") + print(" (attach_plant(plant) is the alternative — it sets both callables)\n") + + +# ─── 4. host-supplied perturbations (the lowering contract) ──────────────── + + +def demo_host_supplied_noise(bk): + print("=== 4. Host-supplied epsilon: sampling stays on the host ===") + sampled = build_custom_controller(bk) + x0 = np.array([1.0]) + u_sampled = bk.to_numpy(sampled.compute(x0)) + eps = sampled._last_epsilon # the draw the controller actually used + if eps is None: # the eager path always records its draw + raise RuntimeError("expected compute() to record the sampled perturbations") + + # A lowered kernel does no sampling: the host draws the perturbations and + # feeds them in, so compute(..., epsilon=...) must reproduce the run. The + # layout is (N, K*D_u), sample-major — what the C-ABI input port expects. + fed = build_custom_controller(bk) + u_fed = bk.to_numpy(fed.compute(x0, None, np.ascontiguousarray(eps.reshape(N_SAMPLES, HORIZON * D_U)))) + print(f" sampled u = {u_sampled[0]:+.10f}") + print(f" fed-eps u = {u_fed[0]:+.10f} (identical: same noise in -> same u out)\n") + + +# ─── 5. trace it: the same callables become graph nodes ──────────────────── + + +def demo_trace_and_lower(bk): + print("=== 5. Tracing: the callables become lowered graph nodes ===") + ctrl = build_custom_controller(bk) + ng = trace_node( + ctrl, + input_shapes={"current_state": (D_X,), "target_state": (D_X,), "epsilon": (N_SAMPLES, HORIZON * D_U)}, + state_shapes={"u": (HORIZON, D_U)}, # recurrent: the nominal sequence + ) + print(f" traced graph: {len(ng.graph.nodes)} nodes, state attrs = {ng.state_attrs}") + print(f" graph outputs: {sorted(ng.output_nodes)}") + + x0 = np.array([1.0]) + eps = np.random.default_rng(0).normal(0.0, 1.0, (N_SAMPLES, HORIZON * D_U)) + feeds = { + "current_state": x0, + "target_state": np.zeros(D_X), + "epsilon": eps, + "state_u": np.zeros((HORIZON, D_U)), + } + out = interpret(ng.graph, feeds) # the interpreter is the .so's oracle + + ref = build_custom_controller(bk) + u_ref = bk.to_numpy(ref.compute(x0, None, eps)) + err = np.max(np.abs(out["out"] - u_ref)) + print(f" interpreter vs live numpy: u max err = {err:.2e} (the Zig VM is checked the same way)") + print(f" published diagnostics: costs{np.asarray(out['costs']).shape}\n") + + +# ─── the rules that decide whether a model can lower ─────────────────────── + + +def print_trace_safety_rules(): + print("=== Trace-safety rules for a lowerable dynamics_fn/cost_fn ===") + for i, rule in enumerate( + ( + "use operators and self.bk ops — no raw numpy (np.sum, np.tile, ...)", + "no x[i] indexing — index with bk.slice_ (row blocks) and .T", + "no bk.sum — write sums as contractions ((z * z) @ W)", + "no rank-differing broadcast — (N, D) * (D,) is rejected; use (1, D)", + "stateful attrs must be rebound (self.u = ...), never mutated in place", + "plant dynamics must be batch-capable — a per-sample x[i] loop cannot be traced", + ), + start=1, + ): + print(f" {i}. {rule}") + print("\n attach_plant's dynamics already follows all of these: LTI plants use one") + print(" batched matmul, nonlinear plants evaluate their own dynamics over the batch.") + + +def main(): + bk = NumpyBackend() + print(f"MPPI demo (backend: {type(bk).__name__}, dt={DT}, N={N_SAMPLES}, K={HORIZON})\n") + demo_custom_model(bk) + demo_plant_model(bk) + demo_config_model(bk) + demo_host_supplied_noise(bk) + demo_trace_and_lower(bk) + print_trace_safety_rules() + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/demos/demo_mppi_nonlinear.py b/demos/demo_mppi_nonlinear.py new file mode 100644 index 0000000..6082abf --- /dev/null +++ b/demos/demo_mppi_nonlinear.py @@ -0,0 +1,342 @@ +"""MPPI on a nonlinear plant: the rollout *is* the plant's dynamics, and it lowers. + +MPPI's rollout is the one place a plant model enters the graph. For a linear +plant that is a batched matmul; for a nonlinear plant it is the plant's own +``dynamics`` — one batch-capable method that serves the eager per-sample path, +the finite-difference linearization, and the traced graph. This demo walks that +from development to a compiled kernel, on an ``InvertedPendulum``: + +1. **One method, two ranks** — ``dynamics(state)`` takes a single ``(2,)`` state + *or* a batch ``(N, 2)`` and returns the derivative with the same rank. It is + the only implementation of the physics; there is no second, batched copy to + drift from it. +2. **Close the loop eagerly** — MPPI (via ``attach_plant``) swings the pendulum + up from a 23-degree tilt and holds it. +3. **Why the nonlinear rollout** — the *predicted* horizon is compared with the + plant's real motion, for the plant's nonlinear model and for the linearized + ``(A, B)``. Far from upright the linearization is wrong; the nonlinear model + is not. +4. **Trace it** — the same ``compute`` call becomes a graph: one ``sin`` node of + shape ``(N, 1)`` per horizon step, so the node count is independent of the + sample count. The graph interpreter is checked against live numpy (it is the + ``.so``'s oracle). +5. **The C-ABI contract** — ``epsilon`` (the perturbations) is a free input + port: sampling stays on the host, which is what makes parity checkable. +6. **Deploy it** (``--build``) — lower to a temp graph, compile the Zig VM with + ReleaseFast, ``dlopen`` the ``.so``, and run the *same* closed loop against + the compiled kernel on the *same* noise, tick by tick. + +``--build`` needs ``zig`` on PATH; without it every other section still runs. +No MuJoCo, no torch required. + +Usage: + python -m demos.demo_mppi_nonlinear # eager + traced (default install) + python -m demos.demo_mppi_nonlinear --build # + compile the .so and run it +""" + +from __future__ import annotations + +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np + +from shinro.codegen import interpret +from shinro.codegen.compose import ComposedGraph +from shinro.codegen.lower_zig import lower_zig +from shinro.codegen.oracle import load_so, output_split, pack_arrays, step_so +from shinro.codegen.trace_node import trace_node +from shinro.controllers.mppi import MPPIController +from shinro.plants.inverted_pendulum import InvertedPendulum +from shinro.utils.array_backend import NumpyBackend + +REPO_ROOT = Path(__file__).resolve().parents[1] +RUNTIME = REPO_ROOT / "src" / "shinro" / "runtime" +BUILD_SO = "--build" in sys.argv + +DT = 0.01 +N_SAMPLES = 64 +HORIZON = 10 +STEPS = 600 +TARGET = np.array([0.0, 0.0]) # upright +X0 = np.array([0.4, 0.0]) # ~23 degrees off upright +SIGMA = 0.5 +U_LIMIT = 1.0 + + +# ─── the plant + the controller, wired the production way ────────────────── + + +def build_plant() -> InvertedPendulum: + return InvertedPendulum(mass=0.1, length=0.5, damping=0.0, gravity=9.81, dt=DT, backend=NumpyBackend()) + + +def build_controller(plant: InvertedPendulum, N: int = N_SAMPLES, K: int = HORIZON) -> MPPIController: + """The live (numpy) MPPI — the reference the traced graph and .so are checked against. + + ``attach_plant`` is the production wiring (``ScenarioFactory`` does the same): + it builds the batched adapter and points MPPI's rollout at the plant's own + ``dynamics``. + """ + ctrl = MPPIController( + num_samples=N, + temperature=0.5, + dt=DT, + horizon=K, + noise_sigma=[SIGMA], + u_min=[-U_LIMIT], + u_max=[U_LIMIT], + seed=42, + backend=NumpyBackend(), + ) + ctrl.attach_plant(plant, Q=np.array([20.0, 1.0]), R=np.array([0.1])) + return ctrl + + +def build_graph(N: int = N_SAMPLES, K: int = HORIZON) -> ComposedGraph: + """Trace the standalone MPPI graph: perturbations in, action + costs out. + + MPPI's Gaussian sampling stays on the host, so ``epsilon`` becomes a free + C-ABI input port of shape ``(N, K*D_u)`` and the traced call never touches + the RNG. The nominal plan ``u`` is the recurrent state. + """ + plant = build_plant() + ctrl = build_controller(plant, N=N, K=K) + ng = trace_node( + ctrl, + input_shapes={"current_state": (2,), "target_state": (2,), "epsilon": (N, K)}, + state_shapes={"u": (K, 1)}, + ) + return ComposedGraph( + graph=ng.graph, + inputs=["current_state", "target_state", "epsilon", "state_u"], + outputs=["out", "costs"], + state_inputs=["state_u"], + state_outputs=["state_u"], + ) + + +# ─── 1. one method, two ranks ────────────────────────────────────────────── + + +def demo_dynamics() -> None: + print("=== 1. The plant's dynamics: one batch-capable method ===") + plant = build_plant() + single = plant.dynamics(np.array([0.3, 0.0]), np.array([0.0])) + batch = plant.dynamics(np.array([[0.3, 0.0], [0.6, 0.1], [-0.2, 0.4]]), np.array([[0.0], [0.1], [-0.3]])) + print(" theta_ddot = (g/l)*sin(theta) + tau/(m*l^2) - (b/(m*l^2))*theta_dot") + print(f" dynamics((2,), (1,)) -> {np.round(single, 6)} shape {single.shape}") + print(f" dynamics((3,2), (3,1)) -> shape {batch.shape} (row i is sample i)") + print(f" batch row 0 == single call : {np.allclose(batch[0], single)}") + print(" Same function serves step(), the finite-difference linearization, and the tracer.\n") + + +# ─── 2. close the loop eagerly ───────────────────────────────────────────── + + +def run_eager(steps: int = STEPS, seed: int = 7): + """Drive the pendulum upright with live numpy MPPI; return (states, us, epsilons).""" + plant = build_plant() + ctrl = build_controller(plant) + ctrl.reset() + plant.state = NumpyBackend().array(X0) + rng = np.random.default_rng(seed) + states, us, epsilons = [X0.copy()], [], [] + for _ in range(steps): + eps = rng.normal(0.0, SIGMA, (N_SAMPLES, HORIZON)) + u = np.asarray(ctrl.compute(plant.get_state(), TARGET, eps)).ravel() + epsilons.append(eps) + us.append(u) + plant.step(u) + states.append(np.asarray(plant.get_state()).ravel().copy()) + return np.array(states), np.array(us), epsilons + + +def demo_closed_loop() -> None: + print("=== 2. Eager closed loop: MPPI swings the pendulum upright ===") + states, _, _ = run_eager() + for i in (0, 50, 150, 300, STEPS): + print(f" t = {i * DT:>5.2f} s theta = {states[i][0]:+.5f} rad theta_dot = {states[i][1]:+.5f}") + print(f" |theta| at the end = {abs(states[-1][0]):.2e} rad (target 0)\n") + + +# ─── 3. why the nonlinear rollout ────────────────────────────────────────── + + +def _rollout(fn, x0, us, dt=DT): + """Explicit Euler over a control sequence — the update MPPI's rollout uses.""" + x = np.array(x0, dtype=float) + for u in us: + x = x + dt * fn(x, np.array([u])) + return x + + +def demo_model_error() -> None: + print("=== 3. Why the nonlinear rollout: the linearization is wrong at a tilt ===") + plant = build_plant() + A_d, B_d = (np.asarray(m) for m in plant.get_model()) # Euler-discretized about upright + A_c, B_c = (A_d - np.eye(2)) / DT, B_d / DT # back to continuous time + theta, tau = 0.8, 0.0 # ~46 degrees — far from the linearization point + x0 = np.array([theta, 0.0]) + + f_true = np.asarray(plant.dynamics(x0, np.array([tau]))).ravel() + f_lin = A_c @ x0 + (B_c @ np.array([tau])).ravel() + print(f" at theta = {theta:.2f} rad, tau = 0:") + print(f" true theta_ddot = {f_true[1]:+.4f} ((g/l)*sin(theta))") + print(f" linearized theta_ddot = {f_lin[1]:+.4f} ((g/l)*theta)") + print(f" -> the linearization is off by {abs(f_lin[1] - f_true[1]) / abs(f_true[1]) * 100:.1f}% at this angle") + + steps = 50 # 0.5 s — long enough for the error to compound + us = [tau] * steps + truth = _rollout(lambda x, u: np.asarray(plant.dynamics(x, u)).ravel(), x0, us) + linear = _rollout(lambda x, u: A_c @ x + (B_c @ u).ravel(), x0, us) + print(f" over {steps} steps ({steps * DT:.2f} s) of the same explicit-Euler rollout:") + print(f" nonlinear model theta = {truth[0]:+.4f}") + print(f" linearized model theta = {linear[0]:+.4f} err {abs(linear[0] - truth[0]):.2e}") + print(" The plant's own dynamics is exact by construction; the (A, B) model is not.") + print(" (The plant's step() integrates semi-implicitly; the rollout uses explicit Euler —") + print(" a small, consistent integration difference, not a model error.)\n") + + +# ─── 4. trace it ─────────────────────────────────────────────────────────── + + +def demo_trace() -> ComposedGraph: + print("=== 4. Trace it: the rollout becomes graph nodes ===") + cg = build_graph() + nodes = cg.graph.nodes + sin_ids = [i for i, n in enumerate(nodes) if n.op == "sin"] + print(f" {len(nodes)} nodes | inputs={cg.inputs}") + print(f" outputs={cg.outputs} | state_outputs={cg.state_outputs}") + print(f" sin nodes at {sin_ids} — one per horizon step, each covering all {N_SAMPLES} samples:") + for i in range(sin_ids[0] - 1, sin_ids[0] + 6): + n = nodes[i] + print(f" [{i:3d}] {n.op:<10} shape={str(n.shape):<8} inputs={list(n.inputs)}") + + # The oracle the .so is checked against: interpreter vs live numpy. + ref = build_controller(build_plant()) + rng = np.random.default_rng(11) + max_err = 0.0 + for _ in range(10): + ref.reset() + x0 = rng.normal(0.0, 0.3, 2) + eps = rng.normal(0.0, SIGMA, (N_SAMPLES, HORIZON)) + feeds = {"current_state": x0, "target_state": TARGET, "epsilon": eps, "state_u": np.zeros((HORIZON, 1))} + out = interpret(cg.graph, feeds) + u_ref = np.asarray(ref.compute(x0, TARGET, eps)).ravel() + max_err = max(max_err, np.max(np.abs(out["out"] - u_ref)).item()) + print(f" interpret() vs live numpy over 10 seeded draws: max |du| = {max_err:.2e} (tier 1e-11)") + + # The structural claim: the batch lives in the shapes, not in the node count. + small, large = build_graph(N=16, K=HORIZON), build_graph(N=64, K=HORIZON) + print(f" node count: N=16 -> {len(small.graph.nodes)}, N=64 -> {len(large.graph.nodes)} (identical)") + print(" only the tensor shapes (and the stack buffer) grow with N.\n") + return cg + + +# ─── 5. the C-ABI contract ───────────────────────────────────────────────── + + +def demo_abi(cg: ComposedGraph) -> None: + print("=== 5. What the host packs (C-ABI) ===") + shapes = (cg.inputs, cg.outputs, cg.state_outputs) + port_nodes = {n.attrs["name"]: n.shape for n in cg.graph.nodes if n.op in ("input", "output")} + for label, names in zip(("in ", "out", "st "), shapes): + for name in names: + if name in port_nodes: + print(f" [{label}] {name:<14} {port_nodes[name]}") + print(f" epsilon = (N, K*D_u) = ({N_SAMPLES}, {HORIZON * 1}) — drawn by the host every tick,") + print(" so the kernel contains no RNG and the same noise can be replayed for parity.\n") + + +# ─── 6. deploy it: the same graph, compiled ──────────────────────────────── + + +def demo_deploy(cg: ComposedGraph) -> None: + print("=== 6. Deploy it: lower, compile the Zig VM, dlopen, run the same loop ===") + if shutil.which("zig") is None: + print(" zig not on PATH — skipping. The trace above is the same graph the .so runs.") + print(" Install zig and re-run `python -m demos.demo_mppi_nonlinear --build`.\n") + return + + # Lower to a TEMP graph path: never clobber the shipped + # src/shinro/runtime/graph_data.zig (the KF+LQR base graph). + tmp = Path(tempfile.mkdtemp(prefix="shinro-mppi-nonlinear-demo-")) + graph_path = tmp / "graph_data.zig" + lower_zig(cg, str(graph_path)) + print(f" lowered to {graph_path} (temp; shipped graph untouched)") + + build = subprocess.run( + [ + "zig", + "build", + "--build-file", + str(RUNTIME / "build.zig"), + "--prefix", + str(tmp), + f"-Dgraph={graph_path}", + "-Doptimize=ReleaseFast", # production mode; Debug is unstripped and huge + ], + capture_output=True, + text=True, + ) + if build.returncode != 0: + print(f" zig build failed:\n{build.stderr.strip()[:400]}") + return + + lib = load_so(tmp) + so_bytes = (tmp / "lib" / "libbase.so").stat().st_size + n_out, n_state = output_split(cg) + print(f" compiled libbase.so (ReleaseFast): {so_bytes / 1024:.0f} KiB, " + f"{n_out} output elements, {n_state} state elements") + + # The deployment loop: the host draws epsilon, the kernel does the rollout. + # The SAME epsilon sequence drives the eager reference, tick for tick. + eager_states, eager_us, epsilons = run_eager() + plant = build_plant() + plant.state = NumpyBackend().array(X0) + plan = np.zeros((HORIZON, 1)) + max_du = 0.0 + max_dx = 0.0 + for step, eps in enumerate(epsilons): + feeds = { + "current_state": np.asarray(plant.get_state()).ravel(), + "target_state": TARGET, + "epsilon": eps, + "state_u": plan, + } + out, state = step_so(lib, pack_arrays(cg, feeds), n_out, n_state) + u_kernel = out[:1] + plan = state[: HORIZON * 1].reshape(HORIZON, 1) + max_du = max(max_du, np.max(np.abs(u_kernel - eager_us[step])).item()) + plant.step(u_kernel) + x_kernel = np.asarray(plant.get_state()).ravel() + max_dx = max(max_dx, np.max(np.abs(x_kernel - eager_states[step + 1])).item()) + + print(" closed loop, kernel vs eager numpy on identical noise:") + print(f" max |u_kernel - u_eager| per tick = {max_du:.2e} (bit-parity tier is 1e-12)") + print(f" max |x_kernel - x_eager| per tick = {max_dx:.2e}") + print(f" final state = {np.round(np.asarray(plant.get_state()).ravel(), 6)}") + print(f" artifacts: {tmp}\n") + + +def main() -> None: + print(f"MPPI on a nonlinear plant (InvertedPendulum, N={N_SAMPLES}, K={HORIZON}, dt={DT})\n") + demo_dynamics() + demo_closed_loop() + demo_model_error() + cg = demo_trace() + demo_abi(cg) + if BUILD_SO: + demo_deploy(cg) + else: + print("=== 6. Deploy it ===") + print(" re-run with --build to lower this graph, compile the Zig VM, and run the .so\n") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/demos/demo_smc.py b/demos/demo_smc.py new file mode 100644 index 0000000..bc6ed6a --- /dev/null +++ b/demos/demo_smc.py @@ -0,0 +1,318 @@ +"""SMC demo: a nonlinear plant, defined as f(x)/g(x), deployed exactly as it lowers. + +Sliding Mode Control is the one controller whose runtime inputs are *live plant +evaluations*. There is no estimator and no plant inside the graph: the host +evaluates ``f(x)`` and ``g(x)`` each tick, and the lowered kernel is pure +arithmetic on ``(x, f_x, g_x)``. This demo walks that contract end to end on a +nonlinear control-affine plant: + + x1_dot = x2 + x2_dot = -a*sin(x1) - b*x2 + c*u (pendulum-like, SISO) + +1. **Define the plant as two functions** — ``f(x)`` (drift, ``(n_x,)``) and + ``g(x)`` (input matrix, ``(n_x, n_u)``). The controller never sees them. +2. **Close the loop eagerly** — the controller drives ``s = c^T x`` to zero. +3. **Model mismatch** — tell the controller the wrong ``a``; the switching + gain ``k1`` is what absorbs the error (that is the point of SMC). +4. **Config-driven** — the shipped ``configs/controllers/smc.toml``. +5. **The controllability guard, both ways** — eager numpy raises + ``RuntimeError``; the lowered graph cannot raise, so it emits ``u = 0`` and + a ``healthy = 0`` flag the host acts on in the same tick. +6. **Trace it** — the same ``compute`` call becomes a graph; the graph + interpreter is checked against live numpy (it is the ``.so``'s oracle). +7. **Deploy it** (``--build``) — lower to a temp graph, compile the Zig VM, + ``dlopen`` the ``.so``, and run the *same* nonlinear closed loop against the + compiled kernel, checking every tick against live numpy. This is "SMC + deployed as is": host owns ``f``/``g``, kernel owns the arithmetic. + +``--build`` needs ``zig`` on PATH (the first build takes ~30 s; later builds +reuse the cache and finish in ~1 s); without it every other section still runs. +No MuJoCo, no torch required. + +Scope note: this is single-surface SMC (one row ``c``). An ``n_u > 1`` plant +needs no code change — ``g_x`` shaped ``(n_x, n_u)`` selects the minimum-norm +branch automatically — and a bank of instances covers multi-axis control. See +``TestSmcOracle`` in ``tests/test_zig_lowering.py`` for the ``n_u = 2`` oracle. + +Usage: + python -m demos.demo_smc # eager + traced (default install) + python -m demos.demo_smc --build # + compile the .so and run it +""" + +from __future__ import annotations + +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np + +from shinro.codegen import interpret +from shinro.codegen.compose import ComposedGraph +from shinro.codegen.lower_zig import lower_zig +from shinro.codegen.oracle import load_so, output_split, pack_arrays, step_so +from shinro.codegen.trace_node import trace_node +from shinro.controllers.smc import SlidingModeController + +BUILD_SO = "--build" in sys.argv + +REPO_ROOT = Path(__file__).resolve().parents[1] +RUNTIME = REPO_ROOT / "src" / "shinro" / "runtime" + +DT = 0.001 +STEPS = 4000 + + +# ─── 1. the plant, as two functions the host owns ────────────────────────── + +A_GRAV = 1.0 # gravity gain in the true plant +B_DAMP = 0.5 # damping +C_CTRL = 1.5 # control authority + + +def f(x) -> np.ndarray: + """The TRUE drift f(x): ``(n_x,) -> (n_x,)``. Evaluated host-side.""" + return np.array([x[1], -A_GRAV * np.sin(x[0]) - B_DAMP * x[1]]) + + +def g(x) -> np.ndarray: + """The TRUE input matrix g(x): ``(n_x,) -> (n_x, n_u)``. + + Constant here, but it may depend on ``x`` — the controller treats it as + data either way. A 1-D ``(n_x,)`` column is also accepted. + """ + return np.array([[0.0], [C_CTRL]]) + + +def f_model(x, a_hat: float) -> np.ndarray: + """The drift the CONTROLLER is TOLD about — ``a_hat`` is our lie.""" + return np.array([x[1], -a_hat * np.sin(x[0]) - B_DAMP * x[1]]) + + +def build_smc( + c: tuple[float, ...] = (1.0, 2.0), + k1: float = 3.0, + k2: float = 1.0, + phi: float = 0.1, + smoother: str = "tanh", + alpha: float = 0.0, +) -> SlidingModeController: + """The demo's SMC: ``c=[1, 2]`` (Hurwitz), tanh boundary layer. + + ``c`` defines the surface ``s = c^T x``; ``c=[1, 2]`` means the polynomial + ``1 + 2p`` (root ``p = -0.5``), so on the surface the error decays as + ``x1_dot = -0.5*x1``. Same gain shape as the shipped ``smc.toml``, with a + stronger ``k1`` so the nonlinear demo converges fast. + """ + return SlidingModeController( + c=list(c), k1=k1, k2=k2, phi=phi, smoother=smoother, alpha=alpha + ) + + +def integrate(x: np.ndarray, u, dt: float) -> np.ndarray: + """One Euler step of the TRUE plant (the demo's ground truth).""" + return x + dt * (f(x) + g(x) @ np.asarray(u).ravel()) + + +# ─── 2. eager closed loop: exact model ───────────────────────────────────── + + +def demo_eager_closed_loop() -> None: + print("=== 2. Eager closed loop on the nonlinear plant (exact model) ===") + ctrl = build_smc() + c = np.array([1.0, 2.0]) + x = np.array([1.0, 0.0]) # 1 rad from equilibrium + print(f" {'step':>5} {'x1':>9} {'x2':>9} {'s=c^Tx':>9} {'u':>9}") + for step in range(STEPS): + u = ctrl.compute(x, f(x), g(x)) # <-- the whole contract, one line + x = integrate(x, u, DT) + if step % 400 == 0: + print(f" {step:>5} {x[0]:>9.5f} {x[1]:>9.5f} {(c @ x).item():>9.5f} {u[0].item():>9.5f}") + s = (c @ x).item() + print(f" final x = {np.round(x, 5)}, |s| = {abs(s):.2e} (surface reached)\n") + + +# ─── 3. model mismatch: this is what k1 pays for ─────────────────────────── + + +def demo_model_mismatch() -> None: + print("=== 3. Model mismatch: the switching gain k1 absorbs the error ===") + a_hat = 2.0 # controller believes gravity gain is 2.0; the plant uses 1.0 + # The mismatch enters s_dot as c^T(f_true - f_model), bounded here by + # |2*(A_GRAV - a_hat)*sin(x1)| = 2.0, so k1 must exceed ~2.0. + for k1 in (1.0, 3.0): + ctrl = build_smc(k1=k1) + x = np.array([1.0, 0.0]) + for _ in range(STEPS): + u = ctrl.compute(x, f_model(x, a_hat), g(x)) + x = integrate(x, u, DT) + s = (np.array([1.0, 2.0]) @ x).item() + verdict = "converges" if abs(s) < 0.1 else "stalled on the mismatch" + print(f" k1={k1:>4.1f} (mismatch bound ~2.0): |s| = {abs(s):.4f} -> {verdict}") + print(f" x = {np.round(x, 5)} at k1=3.0\n") + + +# ─── 4. the shipped config, through the factory ──────────────────────────── + + +def demo_config_driven() -> None: + print("=== 4. Config-driven: the shipped configs/controllers/smc.toml ===") + from shinro.factories import ControllerFactory + from shinro.utils.config_resolver import resolve_config_path + + ctrl = ControllerFactory(str(resolve_config_path("configs/controllers/smc.toml"))).create() + u = ctrl.compute(np.array([1.0, 0.0]), np.array([0.0, 0.0]), np.array([[0.0], [1.0]])) + print(f" c = {ctrl.c.tolist()}, k1 = {ctrl.k1}, smoother = {ctrl._smoother_name}") + print(f" compute(x=[1, 0]) -> u = {np.asarray(u).ravel()[0].item():+.5f}") + print(" (deployment re-lowers one kernel per config; a config change is a re-lower)\n") + + +# ─── 5. the controllability guard: eager raises, the graph flags ─────────── + + +def build_smc_graph(n_u: int = 1) -> ComposedGraph: + """Trace standalone SMC: ``(x, f_x, g_x)`` in, ``(out, healthy)`` out. + + There is no estimator and no plant to compose with — ``compose()`` + deliberately has no role for ``f_x``/``g_x`` (see ``compose.py``), so the + graph is traced standalone and lowered directly, with the dynamics terms as + free C-ABI ports. SMC is memoryless: ``state_outputs`` is empty. + """ + smc = build_smc() + ng = trace_node(smc, input_shapes={"x": (2,), "f_x": (2,), "g_x": (2, n_u)}) + return ComposedGraph( + graph=ng.graph, + inputs=["x", "f_x", "g_x"], + outputs=["out", "healthy"], + state_inputs=[], + state_outputs=[], + ) + + +def demo_guard(graph: ComposedGraph) -> None: + print("=== 5. Loss of controllability: eager raises, the graph flags ===") + # c^T g = 1*1 + 2*(-0.5) = 0 exactly + arrays = {"x": np.array([1.0, 0.5]), "f_x": np.array([0.3, -0.2]), "g_x": np.array([[1.0], [-0.5]])} + print(f" c^T g = {(np.array([1.0, 2.0]) @ arrays['g_x'].ravel()).item():.1f} (below controllability_eps)") + + try: + build_smc().compute(arrays["x"], arrays["f_x"], arrays["g_x"]) + raise AssertionError("expected the eager backend to raise") + except RuntimeError as exc: + print(f" eager numpy : RuntimeError({exc})") + + with np.errstate(divide="ignore", invalid="ignore"): + traced = interpret(graph.graph, arrays) + print(f" traced graph: u = {np.asarray(traced['out']).ravel()[0].item():.1f}, " + f"healthy = {np.asarray(traced['healthy']).ravel()[0].item():.1f}") + print(" A compiled kernel cannot raise (a Zig panic across the C ABI aborts the host),") + print(" so the guard becomes data: zero command + a flag for the host's fault policy.") + print(" Zero-command is the kernel's floor, NOT a safety guarantee for an unstable plant.\n") + + +# ─── 6. trace it: the interpreter is the .so's oracle ────────────────────── + + +def demo_trace_and_check(graph: ComposedGraph) -> float: + print("=== 6. Trace it: interpreter vs live numpy (the .so's oracle) ===") + ctrl = build_smc() + rng = np.random.default_rng(11) + max_err = 0.0 + for _ in range(25): + x = rng.normal(0.0, 0.5, (2,)) + arrays = {"x": x, "f_x": f(x), "g_x": g(x)} # host-evaluated, as in deployment + traced = interpret(graph.graph, arrays) + want = np.asarray(ctrl.compute(arrays["x"], arrays["f_x"], arrays["g_x"])).ravel() + max_err = max(max_err, np.abs(np.asarray(traced["out"]).ravel() - want).max().item()) + assert np.asarray(traced["healthy"]).ravel()[0].item() == 1.0 + + n_out, n_state = output_split(graph) + print(f" graph: {len(graph.graph.nodes)} nodes, inputs {graph.inputs}, outputs {graph.outputs}") + print(f" no recurrent state (memoryless): state_outputs = {graph.state_outputs}") + print(f" max |interpreter - live numpy| over 25 samples = {max_err:.2e}") + print(f" C-ABI buffers: {n_out} output f64 ({n_out * 8} B), {n_state} state f64\n") + return max_err + + +# ─── 7. deploy it: same graph, compiled ──────────────────────────────────── + + +def demo_deploy(graph: ComposedGraph) -> None: + print("=== 7. Deploy it: lower, compile the Zig VM, dlopen, run it ===") + if shutil.which("zig") is None: + print(" zig not on PATH — skipping. The trace above is the same graph the .so runs.") + print(" Install zig and re-run `python -m demos.demo_smc --build` to see it compiled.\n") + return + + # Lower to a TEMP graph path: never clobber the shipped + # src/shinro/runtime/graph_data.zig (the KF+LQR base graph). + tmp = Path(tempfile.mkdtemp(prefix="shinro-smc-demo-")) + graph_path = tmp / "graph_data.zig" + lower_zig(graph, str(graph_path)) + print(f" lowered to {graph_path} (temp; shipped graph untouched)") + + build = subprocess.run( + [ + "zig", + "build", + "--build-file", + str(RUNTIME / "build.zig"), + "--prefix", + str(tmp), + f"-Dgraph={graph_path}", + "-Doptimize=ReleaseFast", # the production mode; Debug is ~11 MB unstripped + ], + capture_output=True, + text=True, + ) + if build.returncode != 0: + print(f" zig build failed:\n{build.stderr.strip()[:400]}") + return + + lib = load_so(tmp) + so_bytes = (tmp / "lib" / "libbase.so").stat().st_size + n_out, n_state = output_split(graph) + print(f" compiled libbase.so (ReleaseFast, stripped): {so_bytes / 1024:.0f} KiB, " + f"{n_out} outputs, {n_state} state") + + # The deployment loop: host evaluates f/g, kernel does the arithmetic. + ctrl = build_smc() + x = np.array([1.0, 0.0]) + max_err = 0.0 + for step in range(STEPS): + packed = pack_arrays(graph, {"x": x, "f_x": f(x), "g_x": g(x)}) + out, _state = step_so(lib, packed, n_out, n_state) + u_kernel, healthy = out[0], out[1] + assert healthy == 1.0, "unexpected controllability fault in the demo trajectory" + + u_ref = np.asarray(ctrl.compute(x, f(x), g(x))).ravel()[0].item() + max_err = max(max_err, abs(u_kernel - u_ref).item()) + x = integrate(x, u_kernel, DT) + if step % 2000 == 0: + print(f" step {step:>5}: x = {np.round(x, 5)}, u = {u_kernel:+.5f}, healthy = {healthy:.0f}") + + print(f" final x = {np.round(x, 5)}") + print(f" max |compiled .so - live numpy| per tick = {max_err:.2e} (bit-parity tier is 1e-12)") + print(f" artifacts: {tmp}\n") + + +def main() -> None: + print(f"SMC demo (plant: x2_dot = -{A_GRAV}*sin(x1) - {B_DAMP}*x2 + {C_CTRL}*u, dt={DT})\n") + demo_eager_closed_loop() + demo_model_mismatch() + demo_config_driven() + graph = build_smc_graph() + demo_guard(graph) + demo_trace_and_check(graph) + if BUILD_SO: + demo_deploy(graph) + else: + print("=== 7. Deploy it ===") + print(" re-run with --build to lower this graph, compile the Zig VM, and run it\n") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/docs/codegen.md b/docs/codegen.md index 8949cb3..ca851a9 100644 --- a/docs/codegen.md +++ b/docs/codegen.md @@ -55,7 +55,7 @@ The pipeline lives in `src/shinro/codegen/`; the Zig VM lives in `src/shinro/run ## Module map | Module | Role | -|--------|------| +| -------- | ------ | | `codegen/tracing.py` | `Tracer` (abstract value), `Graph` / `Node` (graph records), shape checking. Operator overloads (`@`, `+`, `-`, `*`, `.T`) record nodes. | | `codegen/trace_backend.py` | `TraceBackend` — a recording `ArrayBackend` that emits nodes for the named `bk.*` methods components call. | | `codegen/trace_node.py` | `trace_node` / `trace_node_with_state` — run one component call under a `TraceBackend` and return a `NodeGraph`. | @@ -126,8 +126,8 @@ The only time `codegen/` needs an edit for a new component is if it uses a new ## The composition pass -`compose(estimator, controller, plant_dims, input_limits)` wires the fixed ABC -dataflow: +`compose(estimator, controller, plant_dims, input_limits, host_inputs)` wires +the fixed ABC dataflow: ``` y (measurement) ──▶ Estimator ──x_hat──▶ Controller ──u──▶ [clip] ──▶ output @@ -143,7 +143,11 @@ state_P (recurrent) ─────▶ Estimator (any state_* port the trace d `state_P`, `state_u_prev` — fed back as inputs next tick). A state attr the trace detected as mutated without a matching pre-injected placeholder raises: that recursion would be silently frozen at its trace-time value (e.g. the - KF's covariance collapsing to a one-step gain). + KF's covariance collapsing to a one-step gain). **Auxiliary outputs** + published with `ArrayBackend.emit_named_output` (diagnostics — MPPI's + per-sample `costs`, SMC's `healthy` flag) are forwarded from both subgraphs + and appended after `u`, so a deployed binary keeps its observability. A name + published by both components is a loud error, not a silent shadow. - **Controller role mapping:** the controller's inputs are mapped by *role* from its `compute()` signature (`_CONTROLLER_INPUT_ROLES` in `compose.py`), not by hardcoded names. Roles: `state` (names like `x0`, `current_state`, @@ -158,6 +162,11 @@ state_P (recurrent) ─────▶ Estimator (any state_* port the trace d - A controller declaring **`u_prev`** (MPC_DeltaU) shares the estimator's previous-control recurrent port — the same value feeds both, and `state_u_prev` closes the loop. + - A controller may also declare **free host inputs** — ports the *host* fills + each tick rather than the estimator (`compose(host_inputs=...)`, declared + by the component's `host_input_shapes()`). MPPI's `epsilon` is the example: + `(N, K*D_u)` of sampled perturbations, appended last so existing port + layouts are unchanged (sampling stays host-side). - Unmapped input names (e.g. SMC's dynamics terms `f_x`/`g_x`, which need a different wiring model) raise at compose time rather than mis-wiring. - **Controller recurrent state:** the same `state_*` mechanism applies to the @@ -176,13 +185,85 @@ state_P (recurrent) ─────▶ Estimator (any state_* port the trace d `(n,1)`. - **`_merge_and_rewire`**: subgraph `input` nodes are placeholders, not copied — consumers are rewired directly to combined-graph source nodes. Subgraph - `output` nodes are markers and are skipped; `compose` declares the combined - outputs itself. + `output` nodes are markers: `out` and every `state_*` are wired explicitly by + `compose`, and any other named output (an `emit_named_output` diagnostic) is + forwarded by `_forward_diagnostics`, appended after `u`. + +### A controller with no dataflow source: SMC (standalone graph) + +SMC's `compute(x, f_x, g_x)` takes **live plant evaluations** — `f_x = f(x)`, +`g_x = g(x)` — not estimator output. The plant is never a graph citizen +(estimator + controller are lowered; the world stays on the host), so there is +no valid `compose` wiring: `f_x`/`g_x` become free C-ABI input ports the host +fills each tick, and the graph is traced standalone by `trace_node` and lowered +directly (no estimator). `g_x` is `(n_x, n_u)`, unlike the `(n_x,)` every other +controller input uses — the caller supplies that shape explicitly rather than +relying on `build_composed_graph`'s role map. + +Two behavioral conventions make it lowerable at all: + +- **Shape-driven branches.** The scalar-vs-multi-input split keys off + `g_x`'s trace-time shape, not on traced values, and both regimes lower: + `n_u == 1` divides directly, while `n_u > 1` uses the closed-form + minimum-norm solution `u = cg^T (cg cg^T)^{-1} num` — the same point + `np.linalg.lstsq` returned, but written as matmul/transpose/div, which are + graph ops (`lstsq`'s SVD is not). Dot products use column-vector form + because the tracer rejects 1D @ 1D. +- **Guards become data.** The controllability denominator going near-zero is + a `RuntimeError` on the live numpy path, but a graph cannot raise (a Zig + panic across the C ABI aborts the host process). The traced path emits the + same condition as nodes: `cond = lt(norm, eps)`, where `norm` is `abs(c^T g)` + for `n_u == 1` and `‖c^T g‖` otherwise (the two agree when `n_u == 1`), then + forces the fail-safe zero command — the scalar branch zeroes `u` with + `where(cond, 0.0, u_raw)`, the multi-input branch zeroes its `(1,1)` scalar + factor before the final matmul (same effect, and `where`'s operands stay the + same shape). It also publishes `healthy = 1 - cond` through + `ArrayBackend.emit_named_output`, an auxiliary output port. `eps` + (`controllability_eps`) is a per-plant deployment design parameter, not a + numerical constant: the law amplifies `1/c^T g`, so command saturation and + chattering arrive long before the `1e-12` arithmetic floor. Zero-command is + the kernel's floor, not a safety guarantee — the host owns the fault policy + for open-loop-unstable plants. The wiring is **not** a per-scenario edge dict — it's the fixed ABC dataflow, the same for every scenario. What's scenario-specific (clip limits, vector dims) comes from the scenario config. +### A nonlinear rollout: batch-capable `dynamics` (MPPI) + +MPPI's rollout is the one place a *plant model* enters the graph: `compute` +calls `dynamics_fn(x_batch, u_batch, dt)` once per horizon step over `N` +samples. For a linear plant that is a batched matmul; for a nonlinear plant it +is the plant's own `dynamics` — and that is what keeps the lowering uniform. + +`Plant.dynamics(state, control, bk=None)` accepts a single `(n_x,)` state or a +batch `(N, n_x)` and returns the derivative with the same rank, evaluating each +coordinate-wise term elementwise on `(N, 1)` columns. One implementation serves +all three callers — the eager per-sample rollout, the finite-difference +linearization, and the traced graph — so there is no second, separately +maintained batched formula to drift from the first. + +Why that matters for lowering: the sample axis lives in the **node shapes**, not +in the graph's structure. A per-sample loop traced into the graph would emit `N` +copies of the body per step (tens of thousands of nodes at production sizes); +the batched form emits one `sin`/`mul`/`stack` node of shape `(N, 1)` per term, +so the node count is independent of `N` — only `K` and `D_u` matter. The batch +is a contiguous leading axis, which is also what lets the backend vectorize. + +Two mechanical requirements the contract places on a plant's `dynamics`: + +- **Route every backend call through the `bk` it is handed** (defaulting to + `self.bk`). `trace_node` swaps only the *traced component's* backend; the + plant's own stays concrete, so `self.bk.sin(tracer)` would evaluate eagerly. + `mppi.attach_plant` resolves the controller's current backend at call time. +- **No scalar indexing.** The tracer has no `__getitem__`, so coordinates come + from the `x.T` + `bk.slice_(j, j + 1)` + `.T` column idiom; + `utils/batching.py` wraps it as `column` and supplies the rank helpers + (`as_batch` / `as_vector` / `control_batch`). + +What does not lower is dynamics with data-dependent per-row branching — that is +control flow, and per the design doctrine it belongs on the host. + ## The interpreter `interpret(graph, inputs)` walks the graph in execution order (the nodes are @@ -209,9 +290,10 @@ def _matmul(node, values, inputs): An unsupported op raises `NotImplementedError` naming the op to add and listing available ops. The current set (from `ops.py`): -`const`, `input`, `output`, `matmul`, `add`, `sub`, `mul`, `neg`, `transpose`, -`inv`, `reshape`, `clip`, `where`, `copy`, `any`, `tanh`, `relu`, `div`, -`exp`, `argmax`, `one_hot`, `slice`. +`const`, `input`, `output`, `matmul`, `add`, `sub`, `mul`, `div`, `ne`, `lt`, +`neg`, `transpose`, `inv`, `reshape`, `clip`, `where`, `copy`, `any`, `stack`, +`tanh`, `relu`, `exp`, `abs`, `sign`, `pow`, `sin`, `cos`, `min`, `argmax`, +`one_hot`, `slice`, `solve_qp`. ## Lowering to Zig (shipped) @@ -347,11 +429,12 @@ actually emitted by the shipped `base_tracking` graph (names follow the Zig enum in `graph_data.zig`; `cst`/`inp`/`out`/`where_op` are the Zig spellings of `const`/`input`/`output`/`where`): -`const`, `input`, `output`, `matmul`, `add`, `sub`, `mul`, `div`, `neg`, -`transpose`, `inv`, `reshape`, `clip`, `where`, `any`, `copy`, `tanh`, `relu`, -`exp`, `argmax`, `one_hot`, `slice`, `sin`, `cos`, `stack`, `solve_qp`. +`const`, `input`, `output`, `matmul`, `add`, `sub`, `mul`, `div`, `ne`, `lt`, +`neg`, `transpose`, `inv`, `reshape`, `clip`, `where`, `any`, `copy`, `tanh`, +`relu`, `exp`, `abs`, `sign`, `pow`, `min`, `argmax`, `one_hot`, `slice`, `sin`, +`cos`, `stack`, `solve_qp`. -Every interpreter op has a VM switch case. `solve_qp` is special: the +Every interpreter op has a VM switch case. `solve_qp` is special: the interpreter handler solves with the Python `osqp` (eps=1e-6), while the VM drives the baked codegen static solver (same problem, same tolerance), so both sides agree within OSQP's tolerance. Adding a *new* interpreter op is a handler diff --git a/docs/testing.md b/docs/testing.md index 3fc342f..26fd421 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -31,7 +31,7 @@ python3 -m pytest tests/ -k "arm" # keyword-filter tests The [`Makefile`](../Makefile) provides short named targets. `make test-` runs exactly one test file. | Target | Runs | -|--------|------| +| -------- | ------ | | `make test` | Full suite, **skips** `test_very_large_horizon_mpc_times_out` and the opt-in markers (`integration`, `mcp`) | | `make test-all` | Full suite, including the slow horizon test (still excludes the opt-in markers) | | `make test-quick` | Unit tests only (controllers, estimators, trajectories, plants, factories, components, array backend, batched adapter, controllability, mcp server) | @@ -54,6 +54,7 @@ The [`Makefile`](../Makefile) provides short named targets. `make test-` r | `make test-zig` | Generate `src/shinro/runtime/graph_data.zig`, build the Zig VM, run `tests/test_zig_lowering.py` (requires `zig` on PATH) | | `make zig-gen` | Serialize the `base_tracking` composed graph to `src/shinro/runtime/graph_data.zig` only | | `make zig-build` | Compile the Zig VM to `build/lib/libbase.so` and stamp the deployment record (`build/lib/libbase.deployment.json`) via `scripts/stamp_deployment.py` (implies `zig-gen`) | +| `make measure-kernels` | Kernel size metrics for a sweep of MPPI rollout shapes: the C-ABI host buffers, the VM stack buffer, and (with `BUILD=1`) the artifact bytes + compile cost. Static by default (no compiler); `DIMS=3x3x6x3,...` where each entry is `D_x x D_u x N x K`; `JSON=` writes the document | | `make lint` | `ruff check .` + `pyright` on source dirs | The per-file targets run their file unconditionally — `make test-controllers` includes the slow horizon test, unlike `make test` which excludes it. @@ -104,7 +105,7 @@ everywhere (CI or local) unless the runner clears `addopts` explicitly. Shared fixtures are defined in `tests/conftest.py`: | Fixture | Purpose | -|---------|---------| +| --------- | --------- | | `numpy_backend` | A `NumpyBackend` instance | | `torch_backend` | A `TorchBackend` on CPU; skips if `torch` is not installed | | `bk` | Parameterized over `numpy` and `torch`, so any test using `bk` runs twice (once per backend) | diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md new file mode 100644 index 0000000..c63f54c --- /dev/null +++ b/lab-notes/daily/2026-09-14.md @@ -0,0 +1,210 @@ +# Lab Notes — 2026-09-14 + +### 2026-09-14 16:30 UTC — initial + +### 2026-09-14 17:10 UTC — SMC lowered to Zig (standalone graph, where-guarded fail-safe) + +**Goal.** Lower the sliding mode controller into the comptime Zig VM, the way +KF/LQR/PID/MPC already are. SMC is the first controller whose runtime inputs +are **live plant evaluations** (`f_x = f(x)`, `g_x = g(x)`), which forced two +design questions: where the model lives, and what a "cannot happen" guard +means in a kernel that cannot raise. + +**Design decisions.** + +1. **The plant stays on the host.** SMC is traced **standalone** — `(x, f_x, + g_x)` are free C-ABI input ports, `u` is the output. `compose()` deliberately + has no role for `f_x`/`g_x` (`compose.py` refuses them rather than mis-wiring), + and the plant was never a graph citizen anyway (only estimator + controller + are lowered). `g_x` is `(n_x, n_u)` — the caller supplies that shape + explicitly; `build_composed_graph`'s role map assumes `(n_x,)`. +2. **Guards become data, not exceptions.** The near-zero `c^T g` guard was a + `RuntimeError`. A graph cannot raise — the VM is straight-line and a Zig + panic across the C ABI aborts the host process, mid-tick, with the plant + still moving. So the traced path emits `cond = lt(abs(cg), eps)`, + `u = where(cond, 0.0, u_raw)` (fail-safe zero command) and publishes + `healthy = 1 - cond` as an **auxiliary output port**. The numpy path keeps + the raise (developer signal, existing tests unchanged); the two differ only + on the pathological input. +3. **`controllability_eps` is a deployment design parameter, not a numerical + constant.** `1e-12` only protects the division; the law amplifies `1/c^T g`, + so saturation/chattering arrive at `|c^T g| ~ 1e-2`–`1e-4` long before that + floor. The field defaults to `1e-12` (arithmetic-only guard) and is meant to + be set above the smallest `|c^T g|` the plant can legitimately have. A + unit-invariant alignment measure `|c^T g|/(‖c‖‖g‖)` was noted as v2 — it + needs norm reductions + `sqrt`, which the VM does not have. + +**Implementation.** + +- `codegen/ops.py` — new `abs`, `sign`, `pow`, `lt` handlers (`lt` is `ne`'s + ordered sibling; the graph's boolean repr is 1.0/0.0 flags). +- `codegen/tracing.py` — `Tracer.__pow__` (exponent is a config scalar, frozen + as a 0-d const) and `Tracer.__lt__` (dataflow flag, like `__ne__`). +- `codegen/trace_backend.py` — `abs`, `sign`, `ravel` (ravel reuses the reshape + arm, same trick as `hstack`), plus `emit_named_output` → records a named graph + output port. +- `utils/array_backend.py` — concrete no-op `ArrayBackend.emit_named_output`, + so a component can publish a diagnostic without changing its return contract + and without `hasattr`-sniffing the backend. +- `controllers/smc.py` — trace-safe `compute`: `bk.ravel` instead of + `.flatten()`; column-vector form for the contractions (the tracer rejects + 1D @ 1D, which numpy collapses to a scalar); scalar-vs-lstsq split on + `g_x`'s **shape** (trace-time known, not data); `_scalar_control` / + `_vector_control` helpers. `_as_float` gives config coercions an actionable + error (`SMCConfig.k1 must be a number, got 'abc'`). The lstsq branch raises + `NotImplementedError` when traced (no graph op) — closed-form min-norm + `u = cgᵀ(cg cgᵀ)⁻¹ rhs` is the follow-up. +- `runtime/lower.zig` — `.abs` / `.sign` unary arms, `.lt` / `.pow` added to + `BinOp` (`av < bv ? 1 : 0`, `std.math.pow`); `codegen/lower_zig.py` emits the + four new enum entries. +- `tests/test_zig_lowering.py` — `_build_smc_graph`, session fixtures (each + lowering to a **tmp** `graph_path` so the shared `graph_data.zig` is not + clobbered), config variants (`phi=0` → `sign`, `sigmoid` → `abs`, `alpha=0.5` + → fractional `pow`), and `TestSmcOracle`: .so vs interpreter vs live numpy on + 25 seeded samples; `u == 0` + `healthy == 0` on exact-zero `c^T g` while the + live component still raises; a drift guard asserting the graph contains + `abs`/`pow`/`lt`/`where` and the `healthy` port; and an `eps=0.5` build + proving the threshold is baked from config (a usable `|c^T g| = 0.2` trips). +- `tests/test_op_shape_matrix.py` — the new ops join the systematic + (op × shape-class) matrix: `lt` in the binary operand-class loop, and + `abs`/`sign`/positive-base fractional `pow` in the pointwise group. This + covers their broadcasting classes and the zero/±1e3 boundary feeds that the + SMC graph alone does not reach. + +**Results.** `make test`: 1126 passed, 7 skipped (full unit suite). +`make test-zig`: 61 passed, 2 skipped (was 52 collected; +9 SMC +cases). `pytest tests/test_controllers.py`: 130 passed, 2 skipped — numpy +behavior unchanged. `make lint`: ruff clean, pyright 0 errors. Shipped graph +restored with `make zig-gen` (the enum addition does not change +`graph_data.zig` — enum member order does not affect the emitted node table). + +**Caveats.** Zero-command is the kernel's floor, not a safety guarantee: for +an open-loop-unstable plant (inverted pendulum, drone) the host's policy on +`healthy == 0` is the real fail-safe, and the fail-safe is per-plant +(`eps` config). LQR/MPC `.so`s have the same silent-singularity exposure with no +flag at all — SMC's compiled form is now the more diagnosable of the two. + +### 2026-09-14 17:00 UTC — update + +### 2026-09-14 17:30 UTC — update + +### 2026-09-14 23:30 UTC — update + +**What.** Lowered SMC now supports multi-input plants. `_vector_control` no +longer uses `np.linalg.lstsq` — its SVD has no graph op, so the branch raised +`NotImplementedError` when traced. It now computes the same minimum-norm +solution in closed form, `u = cg^T (cg cg^T)^{-1} num`, from +matmul/transpose/divide, all of which the Zig VM already implements: no new op +in `ops.py` or `lower.zig`. + +**Why the closed form is the same point.** `A = c^T g` is a 1 × n_u row, so +`(c^T g) u = num` is one equation in n_u unknowns: `A^-1` does not exist, the +solution set is a hyperplane, and lstsq returns its shortest member. For a +single full-rank row that member is `A^T (A A^T)^{-1} num`. The unit tests +assert equality with `np.linalg.lstsq` numerically and check the reaching law +`c^T f + c^T g u = s_dot_desired` directly. + +**Design decisions.** (1) The guard tests `‖c^T g‖` via `denom**0.5`, not its +square, so `controllability_eps` means the same magnitude as `|c^T g|` on the +scalar branch (identical when n_u == 1). (2) The guard rides the (1,1) scalar +factor, not `u`: `scale = num/denom`, `where(fault, 0, scale)`, then +`cg^T @ scale` — zeroing the scalar zeroes `u`, and `where`'s operands stay +same-shape. The tracer rejects rank-mismatched broadcast (`(1,) / (1,1)` +raises), so `num` is reshaped to `(1,1)`; a throwaway shape spike caught this +before any source edit. (3) The eager vector branch now *raises* `RuntimeError` +on lost controllability (previously it had no check at all), and the guard is +evaluated **before** the divide so the invalid input does not first produce +inf/nan or a numpy divide warning. (4) Both branches share one formula, so +numpy/interpreter/.so compare like-for-like. + +**Files.** `src/shinro/controllers/smc.py` — `_vector_control` rewrite plus +docstrings (class `controllability_eps`, `compute`, the method itself). +`docs/codegen.md` — the SMC section's "multi-input branch raises" bullet +replaced with the closed form and the norm-based guard. + +**Tests.** 3 unit tests in `tests/test_controllers.py`, both backends: closed +form == lstsq min-norm; the command is the shortest exact solution (null-space +vectors built from `[cg[1], -cg[0]]`); multi-input lost-controllability raises. +In `tests/test_zig_lowering.py`, `_build_smc_graph`/`_smc_rand_inputs` gained an +`n_u` parameter, plus a new `smc_multi_so` fixture (tmp `graph_path`, same +discipline as `smc_so`) and 3 `TestSmcOracle` cases: .so vs interpreter vs +numpy on 25 seeded samples (max err < 1e-12, plus the surface-equation check), +fail-safe zero command + `healthy == 0` on `c^T g == [0, 0]`, and a +drift-guard on the pseudo-inverse ops (`transpose` + `matmul`). + +**Results.** `make test`: 1135 passed, 7 skipped. `TestSmcOracle`: 12 passed, +2 skipped (the n_u=2 kernel is built and compared in-process). `make lint`: +ruff clean, pyright 0 errors — `smc.py` clean; the 13 pyright errors in +`test_controllers.py` are pre-existing MPPI-test noise, byte-for-byte the same +count before and after this edit. Shipped graph restored with `make zig-gen` +(content-identical; the manifest was an mtime-only touch). + +**Caveats.** Multi-surface SMC (a surface matrix `C`, m surfaces) remains the +roadmap; this change is its prerequisite — the m = 1, n_u > 1 case of the +general dispatcher. Tall systems (m > n_u) cannot occur under the current +single-surface contract. The `inv` op in `linalg.zig` was considered and +rejected here: a rank-1 system needs only a scalar division, and `la.inv` +panics on an exactly-singular pivot — the C-ABI abort the `healthy` +architecture exists to avoid. Zero-command is still the kernel's floor, not a +safety guarantee; the host owns the fault policy. + +### 2026-09-15 00:00 UTC — update + +**Decision: multi-surface SMC is deferred — single surface remains the shipped +design.** + +We looked at generalizing the controller from one sliding surface to a surface +*matrix* `C` (`m` surfaces, `s = Cx ∈ ℝᵐ`), solving the joint system +`C g(x) u = num` with a trace-time shape dispatch (`m == n_u` → `inv`, +`m < n_u` → `Aᵀ(AAᵀ)⁻¹`, `m > n_u` → reject as unanswerable), and a `det`-based +guard so the m×m `inv` can never be handed a singular matrix — `la.inv` panics +on an exactly-zero pivot, and a Zig panic crossing the C ABI aborts the host, +the one failure mode the `healthy` architecture exists to prevent. The full +design is worked out (see the phase-1 lab note above and this entry); it is +**not** implemented. + +**Why deferred.** + +1. *Field practice is single-surface.* Every SMC implementation we surveyed is + one surface per controller instance (SISO plants), or decentralized — one + instance per axis — for MIMO. Coupled multi-surface SMC is a textbook / + research object, not what production stacks ship. +2. *The standard MIMO answer is a bank of scalar instances.* Each instance gets + its own `c`, its own column of `g`, its own `healthy` flag; nothing new + needs building — phase 1's `n_u`-agnostic kernel already covers each + instance. Quadrotors and per-joint robot loops are controlled this way + (cascaded loops + a mixer), and per-channel flags are more diagnosable than + one aggregated flag. +3. *The failure mode that motivates multi-surface is not fixed by it.* The + min-norm solve is unconstrained and has no notion of per-motor limits, so + actuator saturation — the case where decentralized actually breaks — is an + allocation problem (host priority scheduling, or MPC with the already-lowered + `solve_qp`), not a control-law problem. +4. *Shipping it now is not free.* It is a permanent op-vocabulary expansion + (`det` + Zig kernel + oracle cells), a new config/API surface (`surfaces`), + and guard semantics (`|det|^(1/2m)` thresholds) nobody is asking for. Tested + dormant code is not a liability here, but it *is* a public capability claim + and a maintenance surface. + +**Gate for revisiting (executable, not vibes).** Form `C·g(x)` at the operating +point and read the off-diagonal / diagonal ratio. Below ≈0.1 the decentralized +bank is the correct design; approaching 1 means the channels genuinely couple +and phase 2 is warranted. Other triggers: over-actuated geometry (hexcopter, +tilting rotors) where allocation must live inside the control law; integral +surfaces with deliberate cross-channel terms; or a coupled underactuated plant +where a bank demonstrably fails closed-loop. + +**Prerequisite already landed.** Phase 1 +(`feat/nonlinear-zig-lowering`, commit `6d2ede2`) shipped the `m = 1`, `n_u > 1` +min-norm branch — the `m = 1` case of the deferred dispatcher. Resuming phase 2 +is an extension, not a restart. + +**Docs.** The decision is also recorded in the `smc.py` module docstring +("Scope") and class docstring, so the single-surface choice reads as deliberate +rather than as an unimplemented gap. + +### 2026-09-15 00:30 UTC — update + +### 2026-09-15 01:00 UTC — update + + diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md new file mode 100644 index 0000000..0c44874 --- /dev/null +++ b/lab-notes/daily/2026-09-15.md @@ -0,0 +1,502 @@ +# Lab Notes — 2026-09-15 + +### 2026-09-15 16:18 UTC — initial + +### 2026-09-15 17:20 UTC — MPPI lowering, phases 1–3: `min` op + boundary guard + +**Context.** Start of the MPPI→Zig lowering work (plan: phases 1–8). Phases +1–3 are the op-infrastructure prerequisites; MPPI's own compute rewrite +(phase 4+) is untouched here. + +**Phase 1–2 — the `min` reduction op.** MPPI's softmax needs +``beta = min(costs)`` (the overflow shift); it was the only reduction in the +MPPI path with no matmul identity — every `sum` collapses to a matvec/vecmat +with a ones const, but `min` needs a real kernel. Added end-to-end: + +- `codegen/ops.py` — `@register_op("min")`, `np.min(x, axis)` with `axis` as a + node attr (None → 0-d scalar, 0/1 → an axis of a 2-D input). This is the + interpreter half — the oracle the `.so` is checked against. +- `codegen/trace_backend.py` — `bk.min(x, axis=None)`, emitting `min` with + trace-time-known output shape. +- `runtime/linalg.zig` — `min_all` / `min_axis0` / `min_axis1` (argmax + precedent: the math lives here, the VM only dispatches). Each is NaN-aware + (`isNan(v) or v < best`) to mirror numpy's NaN-wins semantics — a plain + `v < best` silently ignores NaNs and diverges from the oracle on NaN feeds. +- `runtime/lower.zig` — `.min` dispatch arm; the axis rides in the node's + `aux` (0 = None, 1 = axis 0, 2 = axis 1), like `slice`'s start offset. +- `codegen/lower_zig.py` — `min` joins the generated `Op` enum and + `_node_vm_info` maps axis → aux. +- `tests/test_op_shape_matrix.py` — 5 cells in the pointwise group: 1-D full + reduction, 2-D full reduction, 2-D axis 0, 2-D axis 1, and a div-by-zero + NaN-propagation cell. + +**Phase 3 — rank-3 boundary guard: already shipped.** `lower_zig()` already +raises on `ndim > 2` (commit `afe2e2c`, with `TestRankGuard`), so the +`(N,K,D_u)` MPPI epsilon must be flattened to `(N, K·D_u)` at the port +boundary — verified, not re-implemented. Reconfirmed the guard test passes. + +**Results.** `make test`: 1135 passed, 7 skipped (unchanged vs baseline). +`make test-zig`: 64 passed, 2 skipped. Full op/shape matrix: 5 passed (all +`min` cells match numpy to 1e-12, NaN propagation included). `make lint`: +ruff clean, pyright 0 errors. `make zig-gen` refreshed the shipped graph: the +`Op` enum gains `min` (the only `graph_data.zig` change — the shipped KF+LQR +graph has no `min` node). The generated manifest was left at HEAD: `make +zig-gen` also rewrites its `provenance` block with this machine's Python/numpy +versions (added by `0bf27b3`, dropped by `e3c38a1`, and re-added by a +generation here), which is unrelated environment churn for a one-op commit. + +**Caveats.** The `min` op is infrastructure only — nothing in a shipped graph +uses it until MPPI's graph is built (phase 6/7). `axis` support beyond +None/0/1 raises. + +### 2026-09-15 18:51 UTC — MPPI lowering, phase 4: `_quad_form` becomes traceable + +**Goal.** Make `BatchedDynamicsAdapter`'s batched quadratic cost trace, so an +LTI MPPI rollout can lower. The adapter's `bk.sum(z*z*W, axis=1)`/`bk.sum(z*(z@W.T), axis=1)` +is the blocker: the adapter holds the **plant's** backend, not the controller's +`TraceBackend`, so `bk.sum` dispatches through numpy/torch and never sees the +traced operand (the LTI `dynamics_fn` was already operator-only, so no change +there). + +**Change.** A sum over an axis is a matmul identity, so both branches become +contractions built from operators (which `Tracer` overloads intercept and lift +the concrete operands from): diagonal `W` is `(z*z) @ W` — `W` itself is the +contraction vector (matvec `(N,D) @ (D,) -> (N,)`); full `W` is +`(z * (z @ W.T)) @ ones`, the row-wise dot product contracted with a ones +vector. No new VM op. + +**Discovery (matters for phase 5).** The first attempt, `(z*z*W) @ ones`, hit +`ShapeMismatchError: cannot broadcast (N, D) and (D,) (different ranks)`: the +**Tracer's `_broadcast_shape` rejects numpy's rank-differing elementwise +broadcast**, even though the VM's `bcast_flat` supports the `(D,)`-against- +`(N,D)` case (raw graphs exercise it in the op/shape matrix). So in traced +component code, a rank-differing elementwise operand must be reshaped to a +same-rank constant or absorbed into a contraction. The diagonal branch sidesteps +it entirely by contracting with `W`; `(z*(z@W.T))` is already same-rank. + +**Scope.** The nonlinear `_integrate` path (per-sample getitem + `torch.vmap`) +is untouched and remains eager-only; the graph build refuses it in phase 6. + +**Verification.** `tests/test_batched_adapter.py` 15 passed / 3 skipped; +eager MPPI (`-k mppi`) 36 passed / 2 skipped; both `_quad_form` branches traced +through `interpret()` and compared to numpy (diag exact 0.0, full 4.4e-16); +`make test` 1135 passed / 7 skipped; `make lint` 0 errors. 4 pyright errors in +`factories/scenario_factory.py` are pre-existing (verified at HEAD via stash) +and outside `make lint`'s scope. + +**Tooling.** Added a project `.pi-lens.json` with `format.enabled: false`: the +deferred pi-lens autoformatter was `zig fmt`-ing every changed Zig file, +including the generated `graph_data.zig` (breaking generator byte-parity) and +all of `linalg.zig` (118 lines of unrelated churn). Restored before the +phase-1/2 commit. + +### 2026-09-15 18:30 UTC — update + +### 2026-09-15 18:44 UTC — update + +### 2026-09-15 18:52 UTC — update + +### 2026-09-15 19:38 UTC — MPPI lowering, phase 5: trace-safe `compute()` + +**Goal.** Rewrite `MPPIController.compute` so one implementation serves eager +numpy, eager torch, and the tracing backend — the prerequisite for lowering +the LTI rollout to Zig (phases 6–7). + +**Sampling stays host-side.** New optional `epsilon` arg, shape +``(N, K*D_u)`` sample-major. ``epsilon is None`` draws from ``self._rng`` (the +unchanged eager path); the traced call receives the perturbation port's Tracer, +so the branch is a trace-time constant and the RNG is never touched while +tracing. Both paths bridge through ``bk.from_numpy``. + +**The body is now all `bk` ops/operators:** + +- The initial state is batched with a *same-rank* broadcast + (``zeros((N,D_x)) + reshape(x0,(1,D_x))``) — the tracer rejects numpy's + rank-differing tile. +- ``eps_in.T`` is taken once (``(K*D_u, N)``); step k's perturbations are the + row block ``[k*D_u:(k+1)*D_u]`` transposed back to ``(N, D_u)``. +- Every sum is a matmul identity: the control penalty contracts with a ones + vector, the weighted update is one ``(1,N) @ (N, K*D_u)``, the softmax + normalizer is ``(1,N) @ ones(N,1)``. The only genuine reduction is + ``beta = min(costs)`` — the phase-1/2 op. +- ``self.u`` is **rebound**, not mutated in place (``detect_state`` keys off + reassignment), and the receding-horizon shift is K row-slices + a stack (the + tracer has no in-place slice assignment). +- The nominal sequence is bridged up front (``u_nominal = bk.from_numpy(self.u)``): + ``numpy + torch`` raises (only ``torch + numpy`` works), so every use goes + through the backend-native copy. +- The tracking reference lives in ``self._x_ref_holder`` (a 1-element list) and + is stored as a ``(1, D_x)`` row: an array-like attr would be mis-detected as + recurrent state, and the adapter's ``(N,D_x) - x_ref`` needs same-rank form. + +**Diagnostics.** ``costs`` is published via ``emit_named_output`` — a graph +output port, like SMC's ``healthy``. ``_last_epsilon``/``_last_costs`` stay +eager-only (assigned only when concrete), so the tracer detects exactly one +state attr: ``u``. + +**Two bugs the trace smoke test caught:** (1) the softmax normalizer +``(1,N) @ (N,1)`` with the weights on *both* sides sums the squares, not ``w`` +— fixed with a ``ones(N,1)`` const (matmul's inner product is the sum); +(2) numpy ε + torch tensors raise *"Concatenation operation is not implemented +for NumPy arrays"* (numpy's ``__add__`` wins over torch's ``__radd__``) — fixed +by bridging both ``eps_in`` and ``u_nominal`` to the backend. + +**Backend additions** (needed by the eager path, not just the tracer): +``ArrayBackend.min`` (numpy ``np.min`` / torch ``torch.min(...).values``), and +``TorchBackend.clip`` now converts numpy bounds via ``as_tensor`` — +``torch.clamp`` rejects numpy array bounds, and the rewrite moved clipping out +of numpy and into backend space. ``_as_float`` (mirroring ``smc.py``) replaced +bare ``float(dt)``/``float(temperature)`` with named errors. + +**Verification.** MPPI 38 passed / 2 skipped (was 36; +1 test on both backends: +``test_mppi_epsilon_param_matches_sampled`` — feeding one controller's draw +reproduces the sampled run). Adapter 15/3. Trace smoke (N=5, K=3, D_x=D_u=3): +``interpret()`` vs eager numpy is **bit-exact** (u err 0.0, state_u err 0.0), +``state_attrs == ['u']``, outputs ``['costs','out','state_u']``, 122 nodes. +``make test`` 1137 passed / 7 skipped; ``make lint`` 0 errors. + +**Phase 6/7 prerequisite discovered.** The clip nodes carry ``(D_u,)`` bounds +against ``(N, D_u)`` operand shapes. numpy/the interpreter broadcast that, but +the Zig lowerer's clip-blob expansion currently accepts only same-size or +scalar bounds — it must learn numpy broadcasting (``np.broadcast_to``) before +these graphs can lower. + +### 2026-09-15 19:38 UTC — update + +### 2026-09-15 19:48 UTC — `demos/demo_mppi.py`: wiring the model, and the lowering contract + +**What.** Added a MuJoCo-free MPPI demo (runs on the default install — no +MuJoCo/torch) that walks the four ways to give MPPI its `dynamics_fn` / +`cost_fn`, since which one you pick decides whether the controller can lower: +constructor injection (hand-written model), attribute injection (the +`from_config` path), `attach_plant(plant, Q=, R=)` (the `ScenarioFactory` +path), and `from_config` + injection (TOML cannot hold a callable). + +**Also demonstrated:** host-supplied `epsilon` (feeding one controller's draw +reproduces the sampled run exactly — the C-ABI contract a lowered kernel +relies on), and an end-to-end trace of the hand-written model through +`trace_node` + `interpret()` vs live numpy (u max err 0.0, `state_attrs == +['u']`, outputs `['costs','out','state_u']`). The demo closes with the +trace-safety rules for a lowerable model (operator/`bk`-op only, no index, no +`bk.sum`, no rank-differing broadcast, rebound state, LTI for lowering). + +**Why the callables matter.** They are *traced through* — inlined into the +graph — so an eager-only model (`bk.sum`, `x[i]`, raw numpy) traces to a +`NotImplementedError` naming the missing op rather than silently lowering +wrong. `attach_plant`'s LTI path satisfies every rule already. + +**Measured.** Section 5's custom integrator (N=200, K=15, D_x=D_u=1) traces to +**388 nodes** — the first hard data point on the comptime-scale question: a +single-axis MPPI graph is already ~8× the shipped KF+LQR graph (~50 nodes), +and the shipped MPPI config is 3-axis. Worth measuring against a real config +in phase 7 before promising deployment sizing. + +**Verification.** `python -m demos.demo_mppi` runs end to end (closed-loop +regulation, plant tracking |x−x_ref|≈0.007, config build, epsilon equivalence, +trace parity 0.0). `ruff check demos/demo_mppi.py` and `make lint` clean. + +### 2026-09-15 19:51 UTC — update + +### 2026-09-15 19:56 UTC — MPPI lowering, phase 6: graph build + interpreter oracle + +**Goal.** Build the MPPI step graph and pin it against live numpy *before* +compiling anything — plus unblock lowering for the clip shape MPPI produces. + +**1. `lower_zig` clip bounds learn numpy broadcasting.** The clip-blob +expansion accepted only same-size or scalar bounds and raised otherwise, but +MPPI clips an ``(N, D_u)`` sample batch with per-channel ``(D_u,)`` control +limits. Now ``np.broadcast_to(bound, node.shape)`` covers scalar, same-shape, +``(1, D)`` row, ``(D,)`` trailing and ``(N, 1)`` column bounds in one path +(the VM indexes its flat blob by output element, so one entry per element is +still required). ``vals.tolist()`` replaced ``float(v)`` (same result; the +linter's ``unchecked-throwing-call-python`` rule flags bare ``float()``; that +rule is narrow — ``open``/``json``/``os``/``shutil``/``int()``/``float()``). + +**2. Shape-matrix cell.** ``clip_vec_broadcast`` in the elementwise group: +``(D,)`` bounds against a ``(2, 3)`` operand, verified through a real ``.so`` +(the old lowerer raised ``ValueError`` on exactly this input). + +**3. `_build_mppi_graph()` + `TestMppiOracle`.** Standalone graph (SMC +pattern): inputs ``current_state``/``target_state``/``epsilon (N, K*D_u)`` + +recurrent ``state_u (K, D_u)``; outputs ``out`` + ``costs``. Built from an LTI +plant via ``attach_plant`` — the production path, so graph and reference +cannot drift. Four interpreter-only tests (no ``.so``): parity vs live numpy +across 10 seeded draws, the ``costs`` diagnostic port, two-tick recurrence +(the graph's own ``state_u`` fed back, compared against sequential live +``compute``), and a structural drift guard (ops + C-ABI port layout). + +**Bug found in the test, not the code.** The 10-draw parity loop reused one +reference controller, whose nominal sequence accumulates across +``compute()`` calls while the graph was fed ``state_u = 0`` every time — they +diverge from draw 2. Fixed with ``ref.reset()`` per iteration. The recurrence +test passing independently confirmed the graph itself is correct. + +**Measured (trace + lower only, no compile).** Nodes are driven by ``K·D_u`` +(the unrolled rollout), *not* by N; the stack buffer is driven by ``N·K`` +(~50 f64 per sample-step): + +| N | K | nodes | buf_len (f64) | stack buffer | +| --- | --- | --- | --- | --- | +| 6 | 3 | 126 | 1,228 | 9.6 KiB | +| 25 | 5 | 180 | 6,868 | 53.7 KiB | +| 50 | 10 | 315 | 25,193 | 196.8 KiB | +| 100 | 15 | 450 | 73,068 | 570.8 KiB | + +The shipped MPPI config (N=200, K=15, D_u=3) extrapolates to ~1.2 MiB of +stack buffer inside ``shinro_step`` — fine on a desktop/Pi (8 MiB thread +stack), a real constraint on a small-stack RTOS. Node count at that size is +~450–1000, so comptime unrolling is not the risk; the buffer is. That is the +Phase-7 measurement to confirm. + +**Verification.** ``TestMppiOracle`` 4 passed; elementwise shape matrix 1 +passed (clip broadcast through a ``.so``); ``make test`` 1141 passed / 7 +skipped (+4); ``make lint`` 0 errors. Shipped generated artifacts untouched +(no clobbering this run). + +### 2026-09-15 20:06 UTC — MPPI lowering, phase 7: the compiled kernel + +**Goal.** Compile the MPPI graph and verify the `.so` three ways — the point +phases 1–6 were building toward. + +**Fixture + tests.** ``mppi_so`` (session fixture, tmp ``graph_path`` — the +shipped ``graph_data.zig`` stays untouched) and two ``TestMppiOracle`` cases: +three-way parity (.so vs interpreter vs live numpy) over 10 seeded draws, and +C-ABI recurrence — the host feeds the kernel's own ``state_out`` straight back +as ``state_u`` and two ticks must match two sequential live ``compute()`` +calls. ``make test-zig`` went 64 → **70 passed**. + +**Measured: compile time and buffer size scale differently.** Compile time and +``buf_len`` follow ``N·K`` (the per-step tensors); node count follows ``K·D_u`` +(the unrolled rollout), *not* N. Zig compile of the generated comptime VM: + +| N | K | nodes | buf_len (f64) | stack buffer | zig compile | +| --- | --- | --- | --- | --- | --- | +| 6 | 3 | 126 | 1,228 | 9.6 KiB | 0.37 s | +| 25 | 5 | 180 | 6,868 | 53.7 KiB | 0.65 s | +| 50 | 10 | 315 | 25,193 | 196.8 KiB | 1.40 s | +| 100 | 15 | 450 | 73,068 | 570.8 KiB | 3.47 s | + +The shipped config (N=200, K=15, D_u=3) extrapolates to **~1.2 MiB** of stack +buffer and ~6–8 s of compile. **Verdict: comptime unrolling is not the risk** +(~450–1000 nodes, seconds to compile); the `buf: [buf_len]f64` stack array is +the thing to know about — fine on a desktop/Pi (8 MiB thread stack), a real +constraint on a small-stack RTOS. + +**Production-scale correctness.** Beyond the small fixture, N=100/K=15 +(450 nodes, 4 in / 103 out / 45 state) was built and compared: **.so vs +interpreter 1.7e-16, .so vs live numpy 1.1e-16** — float-exact at real size, +not just at toy dims. + +**Verification.** ``make test`` 1143 passed / 7 skipped (+2); ``make test-zig`` +70 passed / 2 skipped (+6); ``make lint`` 0 errors. Generated artifacts +untouched (checked `git status src/shinro/runtime/` after the full suite). + +**Addendum — the three "sizes", measured.** They differ by orders of +magnitude, and only one is the deployment memory story: + +| size | N=6 K=3 | N=100 K=15 | N=200 K=15 (shipped) | D_x=6/D_u=6 N=10 K=4 | +| --- | --- | --- | --- | --- | +| C-ABI input buffer | 0.5 KiB | 35.6 KiB | 70.7 KiB | 2.2 KiB | +| C-ABI output buffer | 9 f64 | 103 f64 | 203 f64 | 16 f64 | +| C-ABI state buffer | 9 f64 | 45 f64 | 45 f64 | 24 f64 | +| VM stack buffer (``buf_len``) | 9.9 KiB | 603.8 KiB | 1,134.9 KiB | 37.1 KiB | +| ``.so`` Debug (dev/test default) | 12.1 MiB | 84.9 MiB | — | 17.4 MiB | +| ``.so`` ReleaseFast (production) | 17.0 KiB | 1,288.8 KiB | 2,607.1 KiB | 65.2 KiB | + +- **Host-side buffers are tiny.** The input buffer is dominated by + ``epsilon`` = ``N*K*D_u`` f64; 70.7 KiB at the shipped config. Packing that + per tick is cheap. +- **The stack buffer is the memory story:** ~1.1 MiB at shipped size (the + ``buf: [buf_len]f64`` inside ``shinro_step``). +- **Always build ReleaseFast for deployment.** The test harness builds Debug: + 12–85 MiB with debug info. ReleaseFast strips that (17 KiB small configs to + 2.6 MiB at N=200) — and ``build.zig`` already notes ReleaseSafe is not + validated (it hangs in ``osqp_solve``), so ReleaseFast is the only release + mode. +- **Cost of the shipped size:** ReleaseFast compile at N=200/K=15 took + **181 s** (measured) and produced a 2.6 MiB ``.so``. One-off per config, but + it is minutes, not seconds — Debug was 3.5 s at N=100. The ``.so`` grows with + N because the VM is a fully unrolled straight-line kernel. + +### 2026-09-15 20:42 UTC — kernel size metric (`make measure-kernels`) + +**Why.** The size numbers from phase 7 were measured ad hoc. They are now a +repeatable metric: ``shinro/codegen/measure.py`` (+ a thin +``scripts/measure_kernels.py`` shim, ``make measure-kernels``), reporting the +three sizes that matter and the compile cost, as a table and optionally a JSON +document. + +- ``graph_metrics(cg, workdir)`` — static derivation from the graph: it lowers + and reads the manifest ``lower_zig`` already writes, so the byte fields + (``buf_bytes``, ``const_blob_bytes``, ``input_bytes``, ``output_bytes``, + ``state_bytes``, per-port detail) come from the same node table the VM + compiles. No compiler, ~0.01 s even at N=200. +- ``kernel_metrics(cg, workdir, optimize)`` — lowers + ``zig build``, records + ``so_bytes`` and ``compile_seconds``. **Fails loudly** when ``zig`` is absent + (the deployment-tool precedent: a metric that silently reported "no + artifact" would be worse than useless). +- ``mppi_lti_graph(D_x, D_u, N, K)`` — builds the standard MPPI sampling graph + for **any** dims via trace-safe injected callables (the plant path fixes + dims to the plant's model). This is what makes the multi-input sweep a + measurement rather than a one-off: `--dims 3x3x6x3,6x6x10x4,...`. +- ``main`` — CLI with ``--dims``/``--optimize``/``--build``/``--json``. Static + by default (fast); ``--build`` compiles (minutes at production sizes). + +**Deliberately not a CI gate.** Unlike the build/deployment manifests (which +are deterministic audit records with a master hash), this document includes a +wall-clock compile time, so it is a measurement *sample*, not a diffable +record. The artifact measurement is exercised in tests only when ``zig`` is on +PATH. + +**Tests.** ``tests/test_measure.py`` (9): the spec parser's four rejections, +byte-field self-consistency against the manifest, ``epsilon`` dominating the +input buffer (``N*K*D_u`` f64), the state port being the nominal plan, the +sweep recording every config, and (zig-gated) the artifact recording. + +**Docs.** ``make measure-kernels`` added to the AGENTS.md Commands list and the +docs/testing.md Make-targets table, so it is discoverable next to +``make compile``. + +**Measured on the tool itself** (ReleaseFast, small config): 131 nodes, +9.9 KiB VM buffer, 18.9 KiB ``.so``, 0.51 s compile. + +**Verification.** ``tests/test_measure.py`` 9 passed; ``make test`` 1152 passed +/ 7 skipped (+9); ``make lint`` 0 errors. Generated artifacts untouched. + +### 2026-09-15 20:49 UTC — the size metric moves into the graph manifest + +**Why.** The sizes were a *standalone* document. They now also live in +``graph_data_manifest.json``, so **every lowered graph self-reports them** — +the audit trail already carried element counts (``buf_len``, +``const_blob_len``) and now carries the byte figures and the C-ABI port sizes. + +**Schema (``lower_zig._graph_manifest``, all additive — no key removed):** + +- ``buf_bytes`` — the VM stack buffer (``buf: [buf_len]f64``). +- ``const_blob_bytes``, ``clip_blob_len`` / ``clip_blob_bytes`` (the VM bakes + ``lo`` and ``hi``, so the byte figure doubles the element count). +- ``input_bytes`` / ``output_bytes`` / ``state_bytes`` (the three C-ABI + buffers the host owns). +- a ``bytes`` entry on every input/output/state port. + +**Consolidation.** ``measure.graph_metrics`` now *reads* those fields instead +of recomputing them (``_shape_bytes`` and the ``math`` import are gone) — one +source of truth, the same node table the VM compiles. The metric command is +unchanged from the user's side. + +**Shipped manifest regenerated** (``make zig-gen``): the KF+LQR graph's own +metric is ``buf_bytes=2544`` (2.5 KiB), ``const_blob_bytes=1008``, +``input_bytes=168``, ``output_bytes=24``, ``state_bytes=120``, +``clip_blob_len=3``. ``graph_data.zig`` content is unchanged (schema-only +change). The machine-specific ``provenance`` block ``make zig-gen`` adds was +dropped again to keep the diff schema-only (pre-existing drift, noted in the +phase-1/2 entry). + +**Tests.** Two added to ``tests/test_measure.py``: the manifest records the +memory metric (self-consistent, MPPI's clip blob is non-empty), and +``graph_metrics`` is a faithful read of it. + +**Verification.** ``tests/test_measure.py`` 11 passed; ``make test`` 1154 +passed / 7 skipped (+2); ``make test-zig`` 70 passed / 2 skipped; ``make lint`` +0 errors. Only ``graph_data_manifest.json`` changed under ``src/shinro/runtime/``. + +### 2026-09-15 20:54 UTC — update + +### 2026-09-15 21:03 UTC — MPPI is deployable: `make compile` wiring for the epsilon port + +**Goal.** `make compile SCENARIO=...` produced deployable `.so` kernels for +LQR/PID/MPC/MPC_DeltaU but not MPPI: `compose()` refuses unmapped controller +inputs, so the sampled-perturbation port had nowhere to go. Now an MPPI +scenario compiles end to end. + +**Design — a generic "host input" role, not an MPPI special case.** + +- `MPPIController.host_input_shapes()` declares the free ports it leaves for + the host: `{"epsilon": (N, K*D_u)}`. The controller owns the knowledge; the + pipeline stays generic. +- `compose(..., host_inputs=(...))` (new, default `()`) declares each as a + composed-graph input port at the controller's own declared shape — not the + `(n_x,)` role default — and wires it through unchanged. They are appended + **last**, so every existing port layout (and golden manifest) is byte-for-byte + unchanged. Names not declared still raise, so the mis-wiring guard stays loud. +- `build_composed_graph(..., plant=None)` (new) attaches the plant when the + controller supports it (`attach_plant` — the same wiring `ScenarioFactory` + uses, so sim and compile cannot disagree) and merges the controller's host + shapes into the trace contract. +- `gen_scenario` passes the plant it already builds for `[plant] type/config` + scenarios. + +**Scenario.** `tests/integration/scenarios/mppi_compile.toml` — KF + MPPI on +`holonomic_base`, compile-ready (`[plant]` type/config + `[compile]`), distinct +from the MuJoCo sim scenario `mppi_base_tracking.toml` (sim-backed `[plant]` +name, no `[compile]`). + +**End-to-end result.** `make compile SCENARIO=mppi_compile.toml` → 493 nodes; +C-ABI inputs `['y','x_ref','u_prev','state_x_hat','state_P','state_u','epsilon']`, +outputs `['u']`, state outputs `['state_x_hat','state_P','state_u','state_u_prev']`. +**Oracle B (.so vs interpret): 20 random inputs, max abs err 6.4e-15 ✓**; +deployment record stamped and verified. The host now owns a 72,000-byte (200×45) +`epsilon` port — the noise it must draw and pack each tick. + +**Tests (+4).** `test_codegen_compose.py::TestComposeHostInputs` — the port is +declared at the controller's shape and appended last, is not recurrent, that an +*undeclared* name still raises, and that the port is a real graph input +(interpreter needs it fed). `test_codegen_build.py` — `build_composed_graph` +with `mppi_base.toml` + the holonomic plant exposes `epsilon == (200, 45)` with +the estimator ports unchanged ahead of it. + +**Verification.** `make test` 1158 passed / 7 skipped (+4); `make test-zig` +70 passed / 2 skipped; `make lint` 0 errors. `build/` artifacts (gitignored) +were produced by the real pipeline. + +**Not done.** Nonlinear MPPI still cannot lower (LTI rollout only), and the +composed MPPI binary is one graph with the KF — the earlier standalone-graph +note applied to the test harness, not to this path. + +### 2026-09-15 21:14 UTC — update + +### 2026-09-15 21:21 UTC — observability: diagnostics survive composition + +**Why.** A deployed composed kernel exposed only the control vector. `compose` +skipped subgraph output markers and declared its own outputs, so every +`emit_named_output` diagnostic — MPPI's per-sample `costs`, SMC's `healthy` +flag — was silently dropped. A blind autopilot: no health signal, no +per-tick rollout cost to watch for divergence. + +**Fix.** `compose._forward_diagnostics` (called after each subgraph merge) +forwards every named output that is not `out` and not `state_*` — exactly the +`emit_named_output` ports — reusing the same source-node resolution the state +ports use (including the input-placeholder case). They are appended **after** +`u` in `cg.outputs`, so existing port layouts and golden manifests are +byte-for-byte unchanged. A name published by both components raises rather +than silently shadowing (and `u` is pre-registered as taken, so a diagnostic +named `u` is caught too). Estimator diagnostics come before the controller's. + +The oracle needed no change: `compare_ports` iterates `cg.outputs`, so a +forwarded diagnostic is automatically `.so`-vs-interpreter checked. + +**Verified end to end.** `make compile SCENARIO=mppi_compile.toml` now reports +`outputs: ['u', 'costs']`; the built C-ABI is `u (3,) 24 B` + `costs (200,) +1600 B`; oracle B 20 inputs, max abs err **1.137e-13** (the larger figure than +before is the costs port joining the comparison — accumulated sums, same +tolerance); deployment record stamped and verified. LQR output layout is +unchanged (`['u']`). + +**Tests (+5).** `TestComposeDiagnostics`: forwarded after `u` (estimator then +controller), a real interpreter-visible port, no-diagnostics leaves the layout +untouched, duplicate names raise, and a diagnostic named `u` raises. The MPPI +build test now also asserts `cg.outputs == ['u', 'costs']` and that +`interpret` returns costs of shape `(200,)`. + +**Docs.** `docs/codegen.md` corrected — it said subgraph output nodes “are +skipped”; now it documents the `out` / `state_*` explicit wiring plus diagnostic +forwarding, and the `host_inputs` free-port mechanism from the previous batch. + +**Verification.** `make test` 1164 passed / 7 skipped (+5); `make test-zig` +70 passed / 2 skipped; `make lint` 0 errors. + +### 2026-09-15 21:28 UTC — update + +### 2026-09-15 21:30 UTC — update + + diff --git a/lab-notes/daily/2026-09-16.md b/lab-notes/daily/2026-09-16.md new file mode 100644 index 0000000..2c684db --- /dev/null +++ b/lab-notes/daily/2026-09-16.md @@ -0,0 +1,438 @@ +# Lab Notes — 2026-09-16 + +### 2026-09-16 16:49 UTC — `demos/demo_smc.py`: SMC on a nonlinear plant, deployed as-is + +**Why.** SMC is the only controller whose runtime inputs are *live plant +evaluations* — the host evaluates `f(x)`/`g(x)` and the lowered kernel is pure +arithmetic on `(x, f_x, g_x)`. There was no runnable example showing that +contract, and `demos/` had no SMC usage at all (grep-confirmed; only MuJoCo +robot demos and `demo_mppi.py`). Question that prompted it: "can SMC be +reliably lowered for nonlinear systems, provided the host supplies f_x/g_x?" +Answer: yes, plant-agnostic, because the plant never enters the graph. + +**What.** New `demos/demo_smc.py` (runnable as `python -m demos.demo_smc`, no +MuJoCo/torch needed). Plant is a pendulum-like nonlinear control-affine system, +defined as the two functions the host owns: + + x1_dot = x2 + x2_dot = -a*sin(x1) - b*x2 + c*u (a=1, b=0.5, c=1.5) + +Seven sections, each a teaching point: + +1. Plant as `f(x)`/`g(x)` + `f_model(x, a_hat)` (the model the controller is + *told*) — makes the model/plant split explicit. +2. Eager closed loop: `s = c^T x` reaches the surface in ~0.3 s, then the + state slides along it (`x1_dot = -0.5*x1` from `c=[1,2]`). +3. Model mismatch: controller told `a_hat=2.0`, plant runs `a=1.0`, so the + mismatch term in `s_dot` is bounded by `|2*(a - a_hat)*sin(x1)| = 2.0`. + `k1=1.0` stalls (`|s|=0.33`), `k1=3.0` converges (`|s|=0.011`) — k1 must + dominate the mismatch bound, stated as a number rather than a slogan. +4. Config-driven: shipped `configs/controllers/smc.toml` through + `ControllerFactory` (`sat`, `phi=0.1`). +5. The controllability guard, both ways: `c^T g = 0` exactly makes eager numpy + raise `RuntimeError`, while the traced graph emits `u=0, healthy=0`. A + compiled kernel cannot raise (Zig panic across the C ABI aborts the host). +6. Trace standalone → `interpret` vs live numpy, bit-exact (max err 0.0 over 25 + seeded host-realistic samples); 37-node graph, inputs `[x, f_x, g_x]`, + outputs `[out, healthy]`, no recurrent state. +7. `--build`: lower to a **temp** graph path, `zig build -Doptimize=ReleaseFast`, + `dlopen`, and run the *same* nonlinear loop against the kernel — host + evaluates f/g each tick, kernel does the arithmetic. + +**Design decisions.** + +- Trace is standalone (`trace_node(smc, {"x": (2,), "f_x": (2,), "g_x": (2,1)})` + plus a hand-built `ComposedGraph`), mirroring the `_build_smc_graph` test + fixture. `compose()` has no role for `f_x`/`g_x`, so there is nothing to + compose with — this is the deployment shape, not a shortcut. +- Lowering goes to `tempfile.mkdtemp()` and `-Dgraph=`: the shipped + `src/shinro/runtime/graph_data.zig` (tracked) is never touched — verified + `git status src/shinro/runtime/` clean after a `--build` run. +- `ReleaseFast` rather than Debug in section 7 because that is the production + mode: **4 KiB stripped** vs ~11 MB unstripped Debug. The deployment artifact + is genuinely tiny for a memoryless controller. +- `--build` is opt-in (default install must run; repo convention from + `demo_mppi.py`), and degrades to a printed hint when `zig` is absent. + +**Gotcha found.** `zig build ... -Dgraph ` (two argv elements) fails with +`Expected -Dgraph to be a string, but received a flag`; it must be a single +`-Dgraph=` element, matching `tests/test_zig_lowering.py`'s `cmd += +[f"-Dgraph={graph_path}"]`. Cost ~10 min of head-scratching for a one-character +fix. + +**Verification.** `python -m demos.demo_smc` — all sections run; `--build` — +ReleaseFast `.so` (4 KiB, 2 outputs, 0 state), closed-loop **max +|.so − live numpy| per tick = 3.33e-16** (bit-parity tier 1e-12), final state +identical to the eager run `[0.14521, -0.07261]`. `pytest +tests/test_controllers.py -k "SMC or smc"` 50 passed. `make lint` 0 errors. +README demo list gained the two `demo_smc` invocations. No `src/` change. + +**Scope note left in the demo.** Single-surface SMC (one row `c`); an +`n_u > 1` plant needs no code change (g_x shape selects the minimum-norm +branch) and multi-axis is a bank of instances. n_u=2 remains covered by +`TestSmcOracle` only — multi-surface (`C` matrix, m surfaces solved jointly) +is still unimplemented, per the module docstring's scope block. + +### 2026-09-16 16:56 UTC — update + +### 2026-09-16 17:38 UTC — nonlinear MPPI rollout, step 1: the batched dynamics path + +**Context.** The MPPI lowering thread (phases 1–7, plus the size metric) ships +the *LTI* rollout: `dynamics_fn` is one batched matmul, so the graph's node +count is independent of the sample count. The follow-up is the nonlinear +rollout — the deferred backlog item framed as "host-side vs VM map primitive; +nodes ∝ N·K". Step 1 is the adapter plumbing that decides the fork. + +**The design call (and why).** Neither of the two framings wins: + +- **Per-sample getitem + Python loop.** The eager nonlinear path loops + `for i in range(N): plant.dynamics(x[i], u[i])`. Traced, each `i` emits its + own ~10-node subgraph: `N·K·body` nodes, ~30,000 at the shipped N=200/K=15. + Yesterday measured ~0.2 s/node ReleaseFast → **~100 min compile**. Dead. +- **VM `map` primitive** (per-row runtime loop over a comptime body) works but + is a real VM/interpreter/lower_zig extension — body-region encoding, nested + execution in the interpreter, a runtime loop inside the comptime `inline + for`. Weeks of work for bodies that are mostly algebraic. +- **Batched nonlinear expression** (chosen). Any coordinate-wise algebraic + `f(x, u)` — sin/cos/products/quotients, i.e. InvertedPendulum, CartPole, + DoublePendulum — is a fixed formula applied to every sample, so it is + expressible with *elementwise ops on `(N, 1)` columns*: extract columns with + the established `x.T` + `bk.slice_(j, j+1)` + `.T` idiom, then `sin`/`mul`/ + `div`/`add`, reassemble with `stack`. Every op is already in the VM (sin/cos + since issue #13). The batch lives in the **node shape**, not the graph + topology → node count independent of N, ~950 at shipped size, the same + scaling as the LTI kernel. The scalar-indexing blocker (`Tracer` not + subscriptable) is **bypassed, not fixed**: `x[i]` is never called on the + graph path; the per-sample loop stays as the oracle. +- Data-dependent per-row branching stays out — per the "binary does raw maths" + doctrine it belongs on the host; the VM `map` primitive remains the + explicitly deferred fallback for genuinely un-batchable bodies. + +**What changed (step 1 = adapter plumbing, no plant yet).** + +- `utils/batched_adapter.py` — path dispatch at construction, most capable + first: `plant.batched_dynamics` (new, trace-safe) → `plant.dynamics` + per-sample Euler (existing; now documented as *not* trace-safe, i.e. the + oracle) → LTI matmul. Duck-typed like `dynamics` (the Plant ABC declares + neither), exposed via a new read-only `dynamics_path` property. The detection + probe for `plant.dynamics` is skipped entirely when the batched surface is + present. +- **The routing trap this step exists to solve.** `trace_node` swaps only the + *traced component's* `self.bk` (`trace_node.py:149`) — the plant's backend + stays concrete, so a plant's `self.bk.sin(tracer)` would compute eagerly and + fail. `dynamics_fn` now takes an optional `bk`, and `mppi.attach_plant` binds + `lambda x, u, dt: adapter.dynamics_fn(x, u, dt, bk=self.bk)`, resolving the + controller's *current* backend at call time (TraceBackend while tracing). + LTI/nonlinear paths ignore `bk`, so existing kernels are unchanged. +- **`mppi.compute` needed zero changes** — the batched body is a drop-in behind + `dynamics_fn(x_batch, v_k)`, same `(N, D_x)` → `(N, D_x)` contract. + +**Tests.** `tests/test_batched_adapter.py::TestBatchedPath` (+4 cells, both +backends): the batched path is selected and matches an independently re-derived +pendulum Euler step; dispatch priority `batched > nonlinear > lti` (fake batched +double, InvertedPendulum, HolonomicMobileRobot); tracers + a caller-supplied +`TraceBackend` emit `sin`/`slice`/`stack` nodes and return a `Tracer`; and +`attach_plant` routes the controller's swapped `bk` into the plant. The double +is a real `Plant` subclass (test-double precedent from `test_components.py`). + +**Verification.** `tests/test_batched_adapter.py` 22 passed / 3 skipped (torch- +only cells); `make test` **1171 passed / 7 skipped** (the LTI MPPI `.so` oracle +cells included); `make lint` **0 errors**; `git status src/shinro/runtime/` +clean (generated graph + manifest untouched); `python -m demos.demo_mppi` — +trace 388 nodes, outputs `['costs','out','state_u']`, interpreter vs live numpy +**u err 0.00e+00** (the lambda refactor is bit-exact). + +**Noted.** `pi-lens` autofixes `README.md` (markdownlint blank line + the +missing `demo_mppi`/`demo_smc` entries the 16:49 entry referenced). An +over-eager `git checkout -- README.md` discarded them; they are restored here. + +**Next (step 2).** `InvertedPendulum.batched_dynamics` — the first plant on the +new path (D_x=2, one sin, ~10 nodes/step), oracle-checked against the +per-sample loop; then the MPPI trace smoke (step 3). + +### 2026-09-16 17:5x UTC — nonlinear MPPI rollout, step 2: InvertedPendulum.batched_dynamics + +**What.** `InvertedPendulum` gains `batched_dynamics(x_batch, u_batch, bk=None)` +— the same `θ̈` formula as `dynamics()`, evaluated on the whole `(N, 2)` batch +with elementwise backend ops. Columns come from the transpose + row-slice +idiom (`x.T`, `bk.slice_(j, j+1)`, `.T`); the derivative is reassembled with +`bk.stack([...]).T`. `dynamics()` keeps its scalar indexing (`state[0]`), which +that idiom is there to avoid. + +**Contract refinement made here (worth having caught early).** The step-1 +adapter called `batched_dynamics(x, u, dt, ...)` and treated the return as the +next state. Better contract, now in force: **`batched_dynamics` returns the +batched derivative `f(x, u)`; the adapter integrates `x + dt·f`.** One +integration site (the adapter), mirroring how `dynamics` is used per-sample, and +the rollout cannot silently disagree with the oracle by integrating differently. +Also corrected the adapter's "semi-implicit Euler" wording — the rollout is +explicit (forward) Euler (`x + dt·f`); the plant's own `step()` is the +semi-implicit one and is never called by the rollout. + +**Consequence (intended).** Every MPPI+InvertedPendulum run now dispatches to +the batched path, so `test_nonlinear_dynamics_matches_euler` silently became a +batched-vs-loop parity test. A dedicated +`test_pendulum_batched_path_matches_per_sample_oracle` makes the claim explicit +(asserts `dynamics_path == "batched"` and parity at `atol=1e-12`), and +`test_pendulum_batched_dynamics_traces` pins that the method emits +`sin`/`slice`/`stack` nodes under a `TraceBackend`. + +**Test-double housekeeping.** `_PerSamplePlant` was added: a plant with +`dynamics` but no batched surface, so `test_dispatch_priority` keeps exercising +the oracle path independently of how far the shipped plants progress (CartPole +and DoublePendulum will become batched in a later step). + +**Verification.** `tests/test_batched_adapter.py` **25 passed / 3 skipped**; +`make test` **1174 passed / 7 skipped**; `make lint` **0 errors**; +`git status src/shinro/runtime/` clean. + +**Next (step 3).** The MPPI trace smoke on the nonlinear plant: `attach_plant` +→ `trace_node` → `interpret()` vs eager numpy, bit-exact, with `detect_state` +finding only `u` and `costs` published. + +### 2026-09-16 18:2x UTC — step 2.5: one batch-capable `dynamics` (collapses `batched_dynamics`) + +**The call (user's).** Step 1–2 had two plant methods: scalar `dynamics` and a +second `batched_dynamics`. The user proposed collapsing them — "N=1 batch is +basically a scalar" — and that is strictly better: it restores the repo's own +doctrine that *the graph is the Python execution, transcribed*. Two +transcriptions of the same physics is exactly the drift risk the framework +tolerates nowhere else. With one method the lowered graph runs the **same +function** the eager path runs. + +**What changed.** + +- `components.Plant.dynamics(state, control, bk=None)` — the ABC docstring now + states the contract: state may be `(n_x,)` **or** `(N, n_x)`, return with + the same rank; `bk` defaults to `self.bk` and must be used for every backend + call (the `trace_node` swap trap). `None` default unchanged for LTI plants. +- **New `utils/batching.py`** — the rank helpers, so the physics body is + written once with no branching: `as_batch` (promote + flag), `as_vector` + (ravel back), `column(bk, x, j)` (the `slice_(x.T, j, j+1).T` idiom in one + readable call), `control_batch` (scalar / `(n_u,)` / `(N, n_u)`). +- **All three nonlinear plants converted** (this had to include CartPole and + DoublePendulum — the adapter has no non-fragile way to tell a scalar-only + `dynamics` apart, so a partial conversion would either break eager MPPI on + them or need a hacky capability probe): + - `InvertedPendulum`, `CartPole` — elementwise columns; `CartPole._compute_accels` + now takes `bk` and is already column-friendly. + - `DoublePendulum` — the interesting one. Its `bk.solve(M, b)` cannot batch + (`M` would be rank-3 `(N,2,2)`, and the graph backend is strictly 2-D), + and the `_make_*` matrix helpers use item assignment, which `Tracer` has + no `__setitem__` for. Replaced by the closed-form 2x2 solve (Cramer's + rule) on `(N,1)` columns. The determinant is + `m2*l1^2*l2^2*(m1 + m2*sin^2(Δ))`, strictly positive for positive masses, + so the division is always safe — no singular pivot. The three `_make_*` + helpers became dead and were removed; their physics moved into the + `dynamics` docstring. + - `step()` in `InvertedPendulum` and `CartPole` now calls `dynamics` instead + of re-typing the acceleration — this removes a *third* copy of the physics + (`DoublePendulum.step` already called it). +- **Adapter collapsed** (`utils/batched_adapter.py`): `dynamics_path` is now + `"nonlinear" | "lti"`; the `"batched"` path, `_integrate`, and the + `torch.vmap`/per-sample-loop machinery are deleted. The nonlinear rollout is + simply `x + dt * plant.dynamics(x_batch, u_batch, bk=...)`. Plants are now + batch-capable by contract, so no auto-vectorization is needed. + +**The oracle survives — and this was the one real trade-off.** Deleting the +scalar implementation deletes the scalar-vs-batched parity test. But that test +only ever caught *rewrite* errors, and with one method there is no rewrite. The +physics stays pinned by the analytic tests in `tests/test_plants.py`, which are +independent references, not self-comparisons: pendulum balancing +(gravity-cancelling torque ⇒ θ̈≈0), DoublePendulum Coriolis against a +hand-built `M` and `C·ω` via `np.linalg.solve`, DoublePendulum balancing, and +the cartpole step tests against hand-computed semi-implicit Euler. + +**New tests.** `tests/test_plants.py::TestBatchCapableDynamics` — every +nonlinear plant returns the right rank for single and batch inputs, a +`(N, n_x)` call equals `N` single-state calls row-for-row, and a scalar control +means "first input, rest zero". This is the coverage guard: a future nonlinear +plant cannot ship scalar-only. In `tests/test_batched_adapter.py`, the +`_BatchedPlant`/`_PerSamplePlant` doubles and the `torch.vmap` class were +deleted (nothing uses vmap now). + +**Verification.** `tests/test_plants.py` 100 passed; `tests/test_batched_adapter.py` + +- `tests/test_plants.py` 113 passed / 1 skipped; `make test` **1174 passed / +5 skipped**; `make lint` **0 errors**; `git status src/shinro/runtime/` clean. + +**Process notes.** `pi-lens` reported stale mid-edit diagnostics twice (it +still cited `inverted_pendulum.py:151` with 3 params after the file on disk had +4; `ast.parse` + a fresh `pyright` run both said clean). Also: an exact-anchor +Python replacement double-inserted retained anchors (a `replace()` that both +included the end anchor in the new text *and* kept it), producing a mangled +adapter; repaired by targeted `str.replace` and re-verified. + +**Next (step 3).** The MPPI trace smoke on the nonlinear plant: +`attach_plant(InvertedPendulum)` → `trace_node` → `interpret()` vs eager numpy, +bit-exact, `detect_state` finding only `u`, `costs` published. + +### 2026-09-16 18:44 UTC — update + +### 2026-09-16 19:3x UTC — step 3: nonlinear MPPI trace smoke (interpreter oracle) + +**Goal.** The first run of the batched plant body *inside the controller graph*: +`attach_plant(InvertedPendulum)` → `trace_node` → `interpret()` vs eager numpy, +bit-exact. The `bk` routing, the rank contract, and the column idiom all have to +survive contact with MPPI's real trace. + +**What landed** (`tests/test_zig_lowering.py`). + +- `_mppi_pendulum_controller` — live MPPI on `InvertedPendulum` via + `attach_plant` (the production path, same as `ScenarioFactory`), D_x = 2, + D_u = 1. Q/R diagonal, `u` clamped to ±0.5. +- `_build_mppi_pendulum_graph(N, K, dt)` — the standalone nonlinear graph + (SMC/`_build_mppi_graph` shape): `current_state`/`target_state`/`epsilon (N, K·D_u)` + in, `out` + `costs` out, `state_u (K, D_u)` recurrent. **Zero changes to + `mppi.compute` or the port contract** — only the plant behind `dynamics_fn` + differs, which is the whole point. +- `TestMppiNonlinearOracle` (5 interpreter-only cases): + 1. `interpret()` == live numpy over 10 seeded draws, `out` and `state_u` both + to 1e-11 (with `ref.reset()` per draw — the nominal plan is stateful). + 2. `costs` still published, shape `(N,)`, finite. + 3. interpreter recurrence: `tick1["state_u"]` fed back reproduces a second + live `compute()`. + 4. drift guard: `sin` is in the op set (the plant's term really is in the + graph), the usual ops are present, `cg.outputs == ["out","costs"]`, + **`cg.state_outputs == ["state_u"]`** — which *is* the "state detection + found only `u`" assertion (the tracking reference must not be promoted to + a recurrent port), and the port shapes are the C-ABI contract. + 5. **the headline property**: tracing the same policy at N=6 and N=24 gives + the *identical node count*, while the `epsilon` port shape differs + ((6, K·D_u) vs (24, K·D_u)). The batch demonstrably lives in the node + shapes, not in the graph topology. That is the measurement that replaces + the ~30,000-node per-sample alternative. + +**Verification.** `TestMppiNonlinearOracle` 5 passed; `make test` **1179 passed / +5 skipped** (+5); `make lint` **0 errors**; `git status src/shinro/runtime/` +clean after the full suite (the `test_zig_lowering` `.so` fixtures did not +clobber the shipped graph). + +**Gotcha re-confirmed (and narrowed).** The shipped `graph_data.zig` + +manifest *did* get overwritten later in the turn — `buf_len` 318 → 120, i.e. a +different base graph, not formatting. It was **not** the test suite: the check +immediately after `make test` was clean, and re-running +`pytest -k MppiNonlinear` twice left the files untouched. The mutation +correlates with the pi-lens turn-end `✓ Zig clean` / `✓ JSON clean` runners +that fire after an edit to a Zig-adjacent file — the same family as the +zig-fmt clobber the repo-root `.pi-lens.json` (`format.enabled=false`) was +meant to stop, so that guard does not cover this runner. Restored with +`git checkout -- src/shinro/runtime/graph_data.zig +src/shinro/runtime/graph_data_manifest.json`. **Rule: check +`git status src/shinro/runtime/` immediately before staging, every time.** + +**Next (step 4).** The `.so`: an `mppi_pendulum_so` fixture (tmp `graph_path`) +plus three-way parity (.so / interpreter / numpy) and C-ABI recurrence, to +confirm the Zig VM handles the nonlinear op mix at real scale. + +### 2026-09-16 19:15 UTC — MPPI nonlinear rollout, steps 4–6: `.so` oracle, scenario compile, docs + +**Step 4 — the compiled kernel.** `mppi_pendulum_so` fixture (session-scoped, +tmp `graph_path`, so the shipped runtime graph is never touched) plus two +`TestMppiNonlinearOracle` cases: + +- `test_so_matches_interpreter_and_numpy` — three-way parity over 10 seeded + draws: `.so` vs interpreter at **1e-13** (the tight tier, including the + `costs` port), `.so` vs live numpy at **1e-11**. +- `test_cabi_recurrence_matches_numpy` — the host feeds the kernel's own + `state_out` back as the next `state_u`; two ticks match two sequential live + `compute()` calls. + +**Step 5 — the e2e scenario compile.** New controller config +`configs/controllers/mppi_inverted_pendulum.toml` (N=200, K=15, D_u=1, torque +clamped ±0.5) and scenario +`tests/integration/scenarios/mppi_pendulum_compile.toml` — the nonlinear twin +of `mppi_compile.toml`, with a real `[plant]` so the pipeline builds the +pendulum, derives the KF's linearized (A, B) via `get_model()`, and wires MPPI +through `attach_plant` (the plant's own batch-capable `dynamics`). `make +compile`: + + wrote build/scenario-mppi-pendulum/graph_data.zig (794 nodes) + inputs: ['y', 'x_ref', 'u_prev', 'state_x_hat', 'state_P', 'state_u', 'epsilon'] + outputs: ['u', 'costs'] + state: ['state_x_hat', 'state_P', 'state_u', 'state_u_prev'] + oracle B (.so vs interpret): 20 random inputs, max abs err 2.132e-14 ✓ + OK: deployment record matches artifacts (master=f7fe5b66…) + +So a nonlinear rollout goes through the *same* pipeline, the same oracle, and +the same deployment stamp as the LTI path — no special-casing anywhere. + +**Step 6 — docs + demo.** `docs/codegen.md`: `min` added to both op lists (it +had been missing since the op landed), and a new subsection under the +composition pass, *“A nonlinear rollout: batch-capable `dynamics` (MPPI)”*, +stating the contract (two mechanical requirements: route through the passed +`bk`; no scalar indexing) and the node-count argument. `demos/demo_mppi.py` +rule 6 no longer claims "nonlinear rollouts stay eager-only" — it now says +plant dynamics must be batch-capable, with the closing note that +`attach_plant` already satisfies everything (LTI via a batched matmul, +nonlinear via the plant's dynamics). + +**Verification.** `make test` **1181 passed / 5 skipped** (+2); `make lint` +**0 errors**; `git status src/shinro/runtime/` clean. `build/` output is +gitignored. + +**Known characteristic (deliberately not optimized now).** Disassembling the +ReleaseFast kernel (`objdump`) shows *partial* SIMD: packed `ymm` doubles are +emitted (`vmulpd` 477, `vsubpd` 329, `vaddpd` 226 at N=64) alongside a large +scalar remainder (`vaddsd` 1643, `vmulsd` 568), with no FMA and no AVX-512; +`sin`/`exp` are libm calls. Scaling N=8→64, packed arithmetic grew 4.4× while +scalar grew 34× — i.e. vectorization *degrades* with N. Likely cause: the VM's +elementwise arms are `inline for` (comptime fully unrolled → straight-line +code), so only LLVM's SLP pass can pack, not the loop vectorizer. Nothing to +fix for correctness — the oracle is bit-exact — but it is headroom: a runtime +row loop (or explicit vector-friendly structure) is the obvious follow-up if +throughput ever becomes the binding constraint. Recorded here so it is not +re-litigated as a surprise. + +### 2026-09-16 19:5x UTC — `demos/demo_mppi_nonlinear.py`: the nonlinear rollout, dev → deployed + +**Why.** The MPPI nonlinear work was verified by tests and a scenario compile, but +there was no runnable narrative. `demo_mppi.py` teaches the wiring and the +lowering contract on an *LTI* plant; this is its nonlinear twin, showing that the +same contract holds when the rollout is the plant's own `dynamics`. + +**Six sections** (`python -m demos.demo_mppi_nonlinear`, no MuJoCo/torch): + +1. **One method, two ranks** — `dynamics((2,))` → `(2,)`, `dynamics((3,2))` → + `(3,2)`, and batch row 0 equals the single call. The same function backs + `step()`, linearization, and the tracer. +2. **Eager closed loop** — MPPI swings the pendulum from 23° to `|theta| = + 1.5e-3` rad in 3 s. +3. **Why the nonlinear rollout** — at θ = 0.46 rad... measured: the linearized + acceleration `(g/l)θ = 15.70` vs the true `(g/l)sin θ = 14.07` — **11.5% off + at 46°**, compounding to a **0.75 rad** θ error over 0.5 s. The plant's own + dynamics is exact by construction. +4. **Trace it** — 515 nodes, one `sin` node per horizon step, each of shape + `(64, 1)`; `interpret()` vs live numpy is **0.00e+00**; node count identical + at N=16 and N=64. +5. **C-ABI** — the four input ports, two outputs, one recurrent state; epsilon is + `(N, K·D_u)`, host-drawn. +6. **Deploy** (`--build`) — ReleaseFast `.so` in a temp dir, `dlopen`, same noise + replayed through both: **max |u_kernel − u_eager| = 4.55e-15**, max + |x_kernel − x_eager| = **2.11e-15** over 600 closed-loop ticks. + +**Two fixes while writing it.** (a) Section 3 originally compared each model's +prediction to the plant's *semi-implicit* `step()`, which mixes integrator error +with model error — and made the linearized model look *better* than the nonlinear +one. Rewritten to compare accelerations directly (the 11.5% number) and to use one +integrator for the horizon comparison, with the semi-implicit/explicit difference +called out as a separate, small effect. (b) The repo's narrow lint rule that +flags `float(...)` outside a `try` (the same one dodged in the phase-6 commit) — +`.item()` on the numpy scalar instead. + +**Verification.** Default run skips section 6 with a hint; `--build` completes. +`make test` **1181 passed / 5 skipped**; `make lint` **0 errors**; `git status +src/shinro/runtime/` clean (the shipped graph is untouched — the demo lowers to a +temp path). README demo list gained the two invocations. + +### 2026-09-16 19:45 UTC — update + + + +### 2026-09-16 19:47 UTC — update + + + +### 2026-09-16 19:47 UTC — update + + diff --git a/scripts/measure_kernels.py b/scripts/measure_kernels.py new file mode 100644 index 0000000..ee0bf29 --- /dev/null +++ b/scripts/measure_kernels.py @@ -0,0 +1,12 @@ +"""Thin shim — kernel size metrics live in ``shinro.codegen.measure``. + +Kept so the metric is runnable as a script (``make measure-kernels``); the +implementation is importable for tooling via ``shinro.codegen.measure``. +""" + +from shinro.codegen.measure import main + +if __name__ == "__main__": + import sys + + sys.exit(main()) diff --git a/src/shinro/codegen/build.py b/src/shinro/codegen/build.py index 18a8ba0..eddbc75 100644 --- a/src/shinro/codegen/build.py +++ b/src/shinro/codegen/build.py @@ -69,6 +69,7 @@ def build_composed_graph( n_x: int, n_u: int, input_limits: tuple | None = None, + plant: object | None = None, ) -> ComposedGraph: """Trace an estimator + controller and compose the closed-loop step graph. @@ -88,6 +89,11 @@ def build_composed_graph( n_u: Plant input dimension. input_limits: Optional ``(lo, hi)`` clip bounds for the controller output, from ``[scenario].input_limits``. + plant: Optional plant instance for controllers whose model comes from + it rather than from their config (MPPI's dynamics/cost). Passed + through ``attach_plant`` when the controller supports it — the + same wiring the simulation path uses, so sim and compile agree. + Ignored by controllers that do not need it (LQR, PID, MPC). Returns: A :class:`ComposedGraph` for one closed-loop step. @@ -95,16 +101,24 @@ def build_composed_graph( est = _factory(EstimatorFactory, estimator_config) ctrl = _factory(ControllerFactory, controller_config) + if plant is not None and hasattr(ctrl, "attach_plant"): + ctrl.attach_plant(plant) + est_input_shapes = {"measurement": (n_x, 1), "control_input": (n_u, 1)} ctrl_input_shapes = { name: (n_u,) if name == "u_prev" else (n_x,) for name in inspect.signature(ctrl.compute).parameters if name != "self" } + # Free host inputs (e.g. MPPI's epsilon) declare their own shapes — the + # role-based defaults above cannot know them. + host_shapes = ctrl.host_input_shapes() if hasattr(ctrl, "host_input_shapes") else {} + ctrl_input_shapes.update(host_shapes) return compose( _trace_with_state(est, est_input_shapes), _trace_with_state(ctrl, ctrl_input_shapes), plant_dims={"n_x": n_x, "n_u": n_u}, input_limits=input_limits, + host_inputs=tuple(host_shapes), ) diff --git a/src/shinro/codegen/compose.py b/src/shinro/codegen/compose.py index 3803511..f542bd5 100644 --- a/src/shinro/codegen/compose.py +++ b/src/shinro/codegen/compose.py @@ -71,6 +71,7 @@ def compose( controller: NodeGraph, plant_dims: dict[str, int], input_limits: tuple[np.ndarray, np.ndarray] | None = None, + host_inputs: tuple[str, ...] = (), ) -> ComposedGraph: """Compose an estimator and controller into one closed-loop step graph. @@ -97,6 +98,12 @@ def compose( detected (attr-diff) must have one — otherwise the recursion would be silently frozen at its trace-time value and this raises instead. + Auxiliary outputs published with :meth:`ArrayBackend.emit_named_output` + (diagnostics such as MPPI's ``costs`` or SMC's ``healthy`` flag) are + forwarded from both subgraphs and appended **after** ``u``, so existing + port layouts are unchanged. Names must be unique across the composed graph + — a silent shadow would hide one component's signal. + Args: estimator: The traced estimator node graph. controller: The traced controller node graph. @@ -106,6 +113,11 @@ def compose( input_limits: Optional ``(lo, hi)`` clip bounds from ``[scenario.input_limits]``. If provided, a ``clip`` node is inserted on the controller output. + host_inputs: Names of controller inputs the *host* fills each tick + rather than the estimator or reference (e.g. MPPI's sampled + perturbations). Each becomes a composed-graph input port whose + shape comes from the traced controller's own declaration; they + are appended last, so existing port layouts are unchanged. Returns: A :class:`ComposedGraph` with the combined graph and port names. @@ -115,6 +127,11 @@ def compose( combined = Graph() + # Auxiliary output ports (diagnostics like MPPI's `costs`) forwarded from + # the subgraphs, and who published each — so a collision is loud. + forwarded_outputs: list[str] = [] + output_owner: dict[str, str] = {"u": "compose (the control output)"} + # --- declare combined-graph inputs first (the merge will reference them) --- y_id = combined.input("y", (n_x,)) x_ref_id = combined.input("x_ref", (n_x,)) @@ -164,6 +181,7 @@ def compose( **state_port_ids, } est_remap, est_source_ids = _merge_and_rewire(combined, estimator.graph, est_input_map) + forwarded_outputs += _forward_diagnostics(combined, estimator, est_remap, est_source_ids, output_owner, "estimator") # The estimator's output (x_hat) — flatten if it's (n,1) to match the # controller's (n,) expectation. The output_nodes value is a subgraph @@ -196,6 +214,8 @@ def compose( # estimator consumes (e.g. MPC_DeltaU's rate input). controller_takes_reference = False state_role_names: list[str] = [] + free_input_names: list[str] = [] + host_names = set(host_inputs) ctrl_input_map: dict[str, int] = {} for node in controller.graph.nodes: if node.op != "input": @@ -204,6 +224,12 @@ def compose( if str(name).startswith("state_"): # Recurrent controller state — wired separately below. continue + if name in host_names: + # A free input the host fills each tick (e.g. MPPI's sampled + # perturbations). Not wired to the estimator or reference: the + # port is declared below from the controller's own shape. + free_input_names.append(name) + continue role = _CONTROLLER_INPUT_ROLES.get(name) if role == "reference": controller_takes_reference = True @@ -215,8 +241,9 @@ def compose( else: raise ValueError( f"controller input '{name}' does not map to a known role " - f"(state / reference / u_prev); extend _CONTROLLER_INPUT_ROLES " - f"in shinro.codegen.compose" + f"(state / reference / u_prev / a declared host input); extend " + f"_CONTROLLER_INPUT_ROLES in shinro.codegen.compose or pass it " + f"in host_inputs" ) if controller_takes_reference or not state_role_names: state_feed_id = x_hat_flat_id @@ -251,7 +278,14 @@ def compose( shape = _lookup_input_shape(controller.graph, port, default=(n_u,)) ctrl_input_map[port] = combined.input(port, shape) + # Free host inputs, declared last so existing port layouts are unchanged. + # Their shapes come from the traced controller's own input placeholders + # (e.g. MPPI's epsilon is (N, K*D_u), not (n_x,)). + for name in free_input_names: + ctrl_input_map[name] = combined.input(name, _lookup_input_shape(controller.graph, name, default=())) + ctrl_remap, ctrl_source_ids = _merge_and_rewire(combined, controller.graph, ctrl_input_map) + forwarded_outputs += _forward_diagnostics(combined, controller, ctrl_remap, ctrl_source_ids, output_owner, "controller") # The controller's output (u) — clip if input_limits provided, then emit. ctrl_out_src_node = controller.graph.nodes[controller.output_nodes["out"]] @@ -337,8 +371,8 @@ def compose( return ComposedGraph( graph=combined, - inputs=["y", "x_ref", "u_prev"] + state_ports + ctrl_state_ports, - outputs=["u"], + inputs=["y", "x_ref", "u_prev"] + state_ports + ctrl_state_ports + free_input_names, + outputs=["u"] + forwarded_outputs, state_inputs=emitted_state_ports + emitted_ctrl_state_ports + ["u_prev"], state_outputs=emitted_state_ports + emitted_ctrl_state_ports + ["state_u_prev"], ) @@ -364,6 +398,57 @@ def compose( } +def _forward_diagnostics( + combined: Graph, + subgraph: NodeGraph, + remap: dict[int, int], + source_ids: dict[str, int], + owner_of: dict[str, str], + owner: str, +) -> list[str]: + """Forward a subgraph's auxiliary named outputs into the combined graph. + + ``out`` (the primary return value) and ``state_*`` (recurrent edges) are + wired explicitly by :func:`compose`. Every *other* named output was + published through :meth:`ArrayBackend.emit_named_output` — a diagnostic + such as MPPI's per-sample ``costs`` or SMC's ``healthy`` flag. Those used + to be dropped (a composed binary exposed only the control vector), leaving + a deployed kernel with no observability. They are appended after ``u``, so + existing port layouts and golden manifests stay unchanged. + + Args: + combined: The combined graph being built. + subgraph: The traced component whose outputs to forward. + remap: Old subgraph node id → combined node id (from the merge). + source_ids: Placeholder input name → combined node id (from the merge). + owner_of: Port name → who published it; updated with the new ports. + owner: This component's label, used in collision errors. + + Returns: + The forwarded port names, in the subgraph's declaration order. + + Raises: + ValueError: If two components publish the same port name (a silent + shadow would hide one of them). + """ + forwarded: list[str] = [] + for name, src_id in subgraph.output_nodes.items(): + if name == "out" or str(name).startswith("state_"): + continue + if name in owner_of: + raise ValueError( + f"output port '{name}' is published by both the {owner_of[name]} " + f"and the {owner}; emit_named_output names must be unique across " + f"a composed graph" + ) + src_node = subgraph.graph.nodes[src_id] + node_id = source_ids[src_node.attrs["name"]] if src_node.op == "input" else remap[src_id] + combined.output(name, node_id) + owner_of[name] = owner + forwarded.append(name) + return forwarded + + def _merge_and_rewire( combined: Graph, subgraph: Graph, diff --git a/src/shinro/codegen/lower_zig.py b/src/shinro/codegen/lower_zig.py index 9b0b644..e970181 100644 --- a/src/shinro/codegen/lower_zig.py +++ b/src/shinro/codegen/lower_zig.py @@ -40,6 +40,10 @@ from shinro.codegen.compose import ComposedGraph from shinro.codegen.tracing import Graph, Node +#: The compiled VM and every C-ABI buffer are f64, so the manifest's byte +#: metrics are element counts times this. +_FLOAT_BYTES = 8 + def lower_zig( cg: ComposedGraph, @@ -101,21 +105,22 @@ def lower_zig( const_blob.extend(float(v) for v in node.attrs["value"].ravel()) elif node.op == "clip": clip_offsets[i] = len(clip_lo) - n_elems = _size(node.shape) bounds: dict[str, list[float]] = {"lo": clip_lo, "hi": clip_hi} for attr, blob in bounds.items(): - vals = np.asarray(node.attrs[attr], dtype=np.float64).ravel() - if vals.size == 1 and n_elems > 1: - # Scalar bound: numpy broadcasts it; the VM's flat blob - # indexing needs one element per output element. - vals = np.full(n_elems, float(vals[0])) - if vals.size != n_elems: + raw = np.asarray(node.attrs[attr], dtype=np.float64) + try: + # numpy broadcasting covers scalar, same-shape, (1, D) row, + # (D,) trailing, and (N, 1) column bounds. The VM indexes + # its flat blob by output element, so one entry per element + # is required — e.g. MPPI clips an (N, D_u) sample batch + # with per-channel (D_u,) control limits. + vals = np.broadcast_to(raw, node.shape).ravel() + except ValueError as exc: raise ValueError( - f"clip node {i}: {attr!r} bound shape {node.attrs[attr].shape} " - f"cannot broadcast against clip shape {node.shape} in the " - f"lowered VM (supported: same-size or scalar bounds)" - ) - blob.extend(float(v) for v in vals) + f"clip node {i}: {attr!r} bound shape {raw.shape} cannot " + f"broadcast against clip shape {node.shape}" + ) from exc + blob.extend(vals.tolist()) # --- output port packing: separate zero-indexed offsets per buffer --- # `outputs` and `state_out` are separate C-ABI buffers, so each needs its @@ -140,6 +145,7 @@ def lower_zig( lines.append(" transpose, inv, reshape, clip, where_op, any,") lines.append(" copy, tanh, relu, exp, argmax, one_hot, slice,") lines.append(" sin, cos, stack, solve_qp,") + lines.append(" abs, sign, pow, lt, min,") lines.append("};") lines.append("") lines.append("pub const Node = struct {") @@ -238,7 +244,7 @@ def _graph_manifest( def _port(name: str) -> dict: for node in g.nodes: if node.op == "output" and node.attrs["name"] == name: - return {"name": name, "shape": list(node.shape)} + return {"name": name, "shape": list(node.shape), "bytes": _size(node.shape) * _FLOAT_BYTES} raise KeyError(f"output port '{name}' not found in graph") solve_qp = None @@ -249,9 +255,7 @@ def _port(name: str) -> dict: nodes = [] for i, node in enumerate(g.nodes): - vm_op, aux = _node_vm_info( - g, i, node, const_offsets, clip_offsets, input_offsets, cg.outputs, cg.state_outputs - ) + vm_op, aux = _node_vm_info(g, i, node, const_offsets, clip_offsets, input_offsets, cg.outputs, cg.state_outputs) rows, cols = _rows_cols(node.shape) nodes.append( { @@ -266,16 +270,37 @@ def _port(name: str) -> dict: } ) + # Size metrics. Element counts above; byte figures here, because the three + # sizes that matter to a deployment differ by orders of magnitude: the + # C-ABI buffers the host packs per tick, the VM's stack buffer + # (`buf: [buf_len]f64`), and the baked blobs. Recorded in the manifest so + # every graph self-reports them (the audit trail already carried buf_len). + input_bytes = sum(_input_size(g, n) for n in cg.inputs) * _FLOAT_BYTES + output_bytes = sum(_output_size(g, n) for n in cg.outputs) * _FLOAT_BYTES + state_bytes = sum(_output_size(g, n) for n in cg.state_outputs) * _FLOAT_BYTES + # clip_blob_len counts one bound's elements (lo == hi in length); the VM + # bakes both, so the byte figure doubles it. + clip_blob_len = sum(_size(n.shape) for n in g.nodes if n.op == "clip") + return { "float_type": "f64", "buf_len": buf_len, + "buf_bytes": buf_len * _FLOAT_BYTES, "const_blob_len": const_blob_len, + "const_blob_bytes": const_blob_len * _FLOAT_BYTES, + "clip_blob_len": clip_blob_len, + "clip_blob_bytes": 2 * clip_blob_len * _FLOAT_BYTES, + "input_bytes": input_bytes, + "output_bytes": output_bytes, + "state_bytes": state_bytes, "has_solve_qp": solve_qp is not None, "nodes_total": len(g.nodes), "nodes": nodes, "ops": sorted(op_histogram), "op_histogram": op_histogram, - "inputs": [{"name": n, "shape": list(_input_shape(g, n))} for n in cg.inputs], + "inputs": [ + {"name": n, "shape": list(_input_shape(g, n)), "bytes": _input_size(g, n) * _FLOAT_BYTES} for n in cg.inputs + ], "outputs": [_port(n) for n in cg.outputs], "state_outputs": [_port(n) for n in cg.state_outputs], "solve_qp": solve_qp, @@ -387,6 +412,12 @@ def _node_vm_info( # slice(x, start, stop): start is the input offset (stop is implicit — # the node's rows*cols is stop - start). return "slice", node.attrs["start"] + if node.op == "min": + # Minimum reduction: the axis rides in aux (0 = None / full, 1 = axis 0 + # down columns, 2 = axis 1 across rows) so the VM's switch selects the + # right reduction loop without an extra node field. + axis = node.attrs.get("axis") + return "min", 0 if axis is None else axis + 1 return node.op, 0 diff --git a/src/shinro/codegen/measure.py b/src/shinro/codegen/measure.py new file mode 100644 index 0000000..826aa33 --- /dev/null +++ b/src/shinro/codegen/measure.py @@ -0,0 +1,252 @@ +"""Kernel size metrics for lowered graphs — a repeatable measurement. + +The three sizes that matter for a deployed kernel differ by orders of +magnitude, so they are reported separately and machine-readably rather than as +a one-off print: + +1. **C-ABI host buffers** — what the host packs/unpacks each tick (the input + buffer is dominated by the ``epsilon`` port for MPPI). +2. **VM stack buffer** — ``buf: [buf_len]f64`` inside ``shinro_step``; this is + the kernel's memory story (a small-stack RTOS cares about it). +3. **Compiled artifact** — the ``.so`` on disk, plus the compile cost. + +``graph_metrics`` is pure derivation from a lowered graph (fast, no compiler). +``kernel_metrics`` adds the build: it lowers, runs ``zig build``, and records +the artifact bytes and the wall-clock compile time. + +Unlike the build/deployment manifests (which are deterministic audit records), +these measurements purposely include a wall-clock compile time, so the JSON is +a measurement sample, not a diffable record — no master hash depends on it. + +Usage: + python -m shinro.codegen.measure # static only + python -m shinro.codegen.measure --build --optimize ReleaseFast + python -m shinro.codegen.measure --dims 3x3x6x3,6x6x10x4 --json out.json + +Each ``--dims`` entry is ``D_x x D_u x N x K`` (the MPPI sampling rollout shape). +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import shutil +import subprocess +import sys +import tempfile +import time + +import numpy as np + +from shinro.codegen.compose import ComposedGraph +from shinro.codegen.lower_zig import lower_zig +from shinro.codegen.runtime_paths import runtime_root +from shinro.codegen.trace_node import trace_node +from shinro.controllers.mppi import MPPIController +from shinro.utils.array_backend import NumpyBackend + +#: Default sweep: the test fixture, a mid-size, and two multi-input systems. +DEFAULT_DIMS = ("3x3x6x3", "3x3x100x15", "6x6x10x4", "8x4x12x5") + + +def parse_dims(spec: str) -> tuple[int, int, int, int]: + """Parse ``D_x x D_u x N x K`` (e.g. ``3x3x100x15``) into four ints.""" + parts = spec.lower().split("x") + if len(parts) != 4: + raise ValueError(f"--dims entry {spec!r} must be D_x x D_u x N x K") + try: + dims = tuple(int(p) for p in parts) + except ValueError as exc: + raise ValueError(f"--dims entry {spec!r} must be four positive integers") from exc + if any(d <= 0 for d in dims): + raise ValueError(f"--dims entry {spec!r} must be positive") + return dims # type: ignore[return-value] + + +def mppi_lti_graph(D_x: int, D_u: int, N: int, K: int, dt: float = 0.05) -> ComposedGraph: + """The standard MPPI sampling graph for an arbitrary LTI system. + + Dynamics/cost are injected as trace-safe operators (``x @ A.T + u @ B.T``, + quadratic forms as contractions) rather than through ``attach_plant``, so + any ``(D_x, D_u)`` can be measured — the plant path fixes the dims to the + plant's model. ``epsilon`` is a free port (sampling stays on the host) and + ``u`` recurs as ``state_u``. + """ + rng = np.random.default_rng(0) + a = np.diag(rng.uniform(0.8, 1.0, D_x)) + b = rng.normal(0.0, 0.15, (D_x, D_u)) + q = np.abs(rng.normal(1.0, 0.2, D_x)) + r = np.abs(rng.normal(0.1, 0.02, D_u)) + + def dynamics(x_batch, u_batch, dt_step): + return x_batch + dt_step * (x_batch @ a.T + u_batch @ b.T) + + def cost(x_batch, u_batch): + return (x_batch * x_batch) @ q + (u_batch * u_batch) @ r + + ctrl = MPPIController( + dynamics_fn=dynamics, + cost_fn=cost, + num_samples=N, + temperature=1.0, + dt=dt, + horizon=K, + noise_sigma=[0.4] * D_u, + u_min=[-1.0] * D_u, + u_max=[1.0] * D_u, + seed=0, + backend=NumpyBackend(), + ) + ng = trace_node( + ctrl, + input_shapes={ + "current_state": (D_x,), + "target_state": (D_x,), + "epsilon": (N, K * D_u), + }, + state_shapes={"u": (K, D_u)}, + ) + return ComposedGraph( + graph=ng.graph, + inputs=["current_state", "target_state", "epsilon", "state_u"], + outputs=["out", "costs"], + state_inputs=["state_u"], + state_outputs=["state_u"], + ) + + +def graph_metrics(cg: ComposedGraph, workdir: pathlib.Path) -> dict: + """Static size metrics for a graph — the manifest ``lower_zig`` writes. + + No compiler, and no duplicated buffer math: ``lower_zig`` records the byte + figures itself (``buf_bytes``, ``const_blob_bytes``, ``clip_blob_bytes``, + and a ``bytes`` entry per port), so this is a faithful read of the same + node table the VM compiles. + """ + graph_path = workdir / "graph_data.zig" + lower_zig(cg, str(graph_path)) + manifest_path = graph_path.with_name("graph_data_manifest.json") + try: + manifest = json.loads(manifest_path.read_text()) + except (OSError, json.JSONDecodeError) as exc: # pragma: no cover - build-time artifact + raise RuntimeError(f"cannot read {manifest_path}: {exc}") from exc + + return { + "nodes_total": manifest["nodes_total"], + "buf_len_f64": manifest["buf_len"], + "buf_bytes": manifest["buf_bytes"], + "const_blob_bytes": manifest["const_blob_bytes"], + "clip_blob_bytes": manifest["clip_blob_bytes"], + "input_bytes": manifest["input_bytes"], + "output_bytes": manifest["output_bytes"], + "state_bytes": manifest["state_bytes"], + "has_solve_qp": manifest["has_solve_qp"], + "ops": manifest["ops"], + "inputs": manifest["inputs"], + "outputs": manifest["outputs"], + "state_outputs": manifest["state_outputs"], + } + + +def kernel_metrics(cg: ComposedGraph, workdir: pathlib.Path, optimize: str) -> dict: + """Build the kernel and record its artifact bytes and compile cost. + + Fails loudly (raises) when ``zig`` is missing: a measurement tool that + silently reported "no artifact" would be worse than useless. + """ + if shutil.which("zig") is None: + raise RuntimeError("zig is not on PATH — cannot measure the compiled artifact") + graph_path = workdir / "graph_data.zig" + lower_zig(cg, str(graph_path)) + cmd = [ + "zig", + "build", + "--build-file", + str(runtime_root() / "build.zig"), + "--prefix", + str(workdir), + f"-Dgraph={graph_path}", + f"-Doptimize={optimize}", + ] + started = time.perf_counter() + result = subprocess.run(cmd, capture_output=True, text=True) + elapsed = time.perf_counter() - started + if result.returncode != 0: + raise RuntimeError(f"zig build failed: {result.stderr.strip()[:400]}") + so_path = workdir / "lib" / "libbase.so" + if not so_path.exists(): + raise RuntimeError(f"zig build produced no libbase.so under {workdir}") + return {"optimize": optimize, "so_bytes": so_path.stat().st_size, "compile_seconds": round(elapsed, 2)} + + +def measure(specs: list[tuple[int, int, int, int]], optimize: str, build: bool) -> dict: + """Measure each ``(D_x, D_u, N, K)`` spec; returns the metrics document.""" + configs = [] + for d_x, d_u, n, k in specs: + label = f"D_x={d_x} D_u={d_u} N={n} K={k}" + with tempfile.TemporaryDirectory(prefix="shinro-measure-") as tmp: + workdir = pathlib.Path(tmp) + cg = mppi_lti_graph(d_x, d_u, n, k) + entry = { + "label": label, + "D_x": d_x, + "D_u": d_u, + "N": n, + "K": k, + "graph": graph_metrics(cg, workdir), + "kernel": kernel_metrics(cg, workdir, optimize) if build else None, + } + configs.append(entry) + print(_row(entry), flush=True) + return {"generated_by": "shinro.codegen.measure", "float_type": "f64", "build": build, "optimize": optimize, "configs": configs} + + +def _row(entry: dict) -> str: + """One table line: the three sizes side by side (see ``_header``).""" + g = entry["graph"] + kernel = entry["kernel"] + so = f"{kernel['so_bytes'] / 1024:.1f}" if kernel else "-" + secs = f"{kernel['compile_seconds']:.2f}" if kernel else "-" + return ( + f"{entry['label']:24s} {g['nodes_total']:6d}" + f" {g['input_bytes'] / 1024:9.1f} {g['buf_bytes'] / 1024:9.1f}" + f" {so:>9s} {secs:>10s}" + ) + + +def _header(optimize: str, build: bool) -> str: + mode = optimize if build else "static only (no build)" + return ( + f"{'config':24s} {'nodes':>6s} {'in KiB':>9s} {'VM KiB':>9s} {'so KiB':>9s} {'compile s':>10s}" + f" [{mode}]" + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Measure lowered MPPI kernel sizes.") + parser.add_argument("--dims", default=",".join(DEFAULT_DIMS), help="comma-separated D_x x D_u x N x K entries") + parser.add_argument("--optimize", default="ReleaseFast", help="zig optimize mode (production: ReleaseFast)") + parser.add_argument("--build", action="store_true", help="compile and measure the artifact (slower)") + parser.add_argument("--json", dest="json_path", help="also write the metrics document here") + args = parser.parse_args(argv) + + try: + specs = [parse_dims(s) for s in args.dims.split(",") if s.strip()] + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + print(_header(args.optimize, args.build)) + document = measure(specs, args.optimize, args.build) + + if args.json_path: + out = pathlib.Path(args.json_path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(document, indent=2, sort_keys=True) + "\n") + print(f"\nwrote {out}") + return 0 + + +if __name__ == "__main__": # pragma: no cover + sys.exit(main()) diff --git a/src/shinro/codegen/ops.py b/src/shinro/codegen/ops.py index 43a603a..a3ba8c4 100644 --- a/src/shinro/codegen/ops.py +++ b/src/shinro/codegen/ops.py @@ -101,6 +101,17 @@ def _ne(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray] return (values[node.inputs[0]] != values[node.inputs[1]]).astype(np.float64) +@register_op("lt") +def _lt(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: + """Elementwise ``a < b`` as 1.0/0.0 floats — the ordered sibling of ``ne``. + + ``ne`` is the graph's only other boolean; it cannot express an ordering, so + threshold/guard conditions (e.g. SMC's near-zero ``c^T g`` guard) need this + op. Consumed by ``where``/``any`` exactly like ``ne``. + """ + return (values[node.inputs[0]] < values[node.inputs[1]]).astype(np.float64) + + @register_op("neg") def _neg(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: return -values[node.inputs[0]] @@ -216,11 +227,41 @@ def _exp(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray return np.exp(values[node.inputs[0]]) +@register_op("abs") +def _abs(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: + return np.abs(values[node.inputs[0]]) + + +@register_op("sign") +def _sign(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: + """Sign function matching ``np.sign``: -1 / 0 / +1 (0 maps to 0).""" + return np.sign(values[node.inputs[0]]) + + +@register_op("pow") +def _pow(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: + return np.power(values[node.inputs[0]], values[node.inputs[1]]) + + @register_op("argmax") def _argmax(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: return np.asarray(np.argmax(values[node.inputs[0]])) +@register_op("min") +def _min(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: + """Minimum reduction, numpy ``min`` semantics. + + ``axis`` is a node attr: ``None`` collapses to a 0-d scalar, ``0``/``1`` + reduce that axis of a 2-D input. Introduced for MPPI's softmax shift + (``beta = min(costs)``, the overflow guard); the axis support keeps the op + a complete reduction rather than a one-off. The Zig mirror is + ``linalg.min_all`` / ``min_axis0`` / ``min_axis1``, dispatched on the + node's ``aux`` (0 = None, 1 = axis 0, 2 = axis 1). + """ + return np.min(values[node.inputs[0]], axis=node.attrs.get("axis")) + + @register_op("one_hot") def _one_hot(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndarray]) -> np.ndarray: # one_hot(x, depth): x is a scalar index → one-hot row vector. diff --git a/src/shinro/codegen/scenario_gen.py b/src/shinro/codegen/scenario_gen.py index bc7a4ed..9ab3054 100644 --- a/src/shinro/codegen/scenario_gen.py +++ b/src/shinro/codegen/scenario_gen.py @@ -199,6 +199,7 @@ def gen_scenario(scenario_path: str, out_dir: str) -> tuple: # (name only) are sim-only and ignored here. Explicit [compile] n_x/n_u and # config A/B win over derived. plant_cfg = spec["plant"] + plant = None if plant_cfg is not None and "type" in plant_cfg and "config" in plant_cfg: with open(resolve_config_path(plant_cfg["config"]), "rb") as f: plant = _PLANT_REGISTRY[plant_cfg["type"]].from_config( @@ -237,6 +238,7 @@ def gen_scenario(scenario_path: str, out_dir: str) -> tuple: n_x, n_u, input_limits=spec["input_limits"], + plant=plant, ) out = Path(out_dir) out.mkdir(parents=True, exist_ok=True) diff --git a/src/shinro/codegen/trace_backend.py b/src/shinro/codegen/trace_backend.py index 90948a7..02b1973 100644 --- a/src/shinro/codegen/trace_backend.py +++ b/src/shinro/codegen/trace_backend.py @@ -157,10 +157,7 @@ def hstack(self, arrays: list[Tracer]) -> Tracer: if not arrays: raise NotImplementedError("TraceBackend.hstack of empty list") if any(len(a.shape) != 1 for a in arrays): - raise NotImplementedError( - "TraceBackend.hstack only supports 1-D arrays " - f"(got shapes {[a.shape for a in arrays]})" - ) + raise NotImplementedError(f"TraceBackend.hstack only supports 1-D arrays (got shapes {[a.shape for a in arrays]})") stacked = self.stack(arrays) n = sum(a.shape[0] for a in arrays) return self._emit("reshape", [stacked], (n,), target_shape=(n,)) @@ -174,6 +171,15 @@ def reshape(self, x: Tracer, *shape: int) -> Tracer: out_shape = tuple(shape) return self._emit("reshape", [x], out_shape, target_shape=out_shape) + def ravel(self, x: Tracer) -> Tracer: + # Flatten to 1-D. Pure data movement — emitted as a reshape to the + # trace-time-known element count, so the VM reuses its reshape arm + # rather than growing a dedicated op (same trick as hstack). + n = 1 + for d in x.shape: + n *= d + return self._emit("reshape", [x], (n,), target_shape=(n,)) + # --- deterministic-policy ops (NN controllers in deterministic mode) --- def tanh(self, x: Tracer) -> Tracer: @@ -196,6 +202,28 @@ def div(self, a: Tracer, b: Tracer) -> Tracer: def exp(self, x: Tracer) -> Tracer: return self._emit("exp", [x], x.shape) + def abs(self, x: Tracer) -> Tracer: + return self._emit("abs", [x], x.shape) + + def sign(self, x: Tracer) -> Tracer: + return self._emit("sign", [x], x.shape) + + def min(self, x: Tracer, axis: int | None = None) -> Tracer: + # Minimum reduction. axis=None collapses to a 0-d scalar (shape ()), + # axis=0/1 reduce that axis of a 2-D input — numpy's semantics, so the + # interpreter handler is a one-liner. The Zig VM carries the axis in + # the node's aux (0 = None, 1 = axis 0, 2 = axis 1) and dispatches to + # linalg.min_all / min_axis0 / min_axis1. + if axis is None: + out_shape: tuple[int, ...] = () + elif axis == 0: + out_shape = x.shape[1:] + elif axis == 1: + out_shape = (x.shape[0],) + x.shape[2:] + else: + raise NotImplementedError(f"TraceBackend.min supports axis in (None, 0, 1), got {axis}") + return self._emit("min", [x], out_shape, axis=axis) + def argmax(self, x: Tracer) -> Tracer: # numpy argmax over the last axis collapses it to a scalar index. return self._emit("argmax", [x], ()) @@ -218,6 +246,13 @@ def to_numpy(self, x: Tracer) -> Tracer: def from_numpy(self, x: Any) -> Tracer: return _lift(self.g, x) + def emit_named_output(self, name: str, value: Any) -> None: + # Record an auxiliary graph output port. This is how a component + # publishes a diagnostic (e.g. SMC's `healthy` controllability flag) + # alongside its primary return value without changing its return + # contract — the eager backends ignore it (see ArrayBackend). + self.g.output(name, _lift(self.g, value).node) + def allclose(self, a: Tracer, b: Tracer) -> bool: # allclose on tracers is a runtime check — not meaningful at trace # time. Components don't call this in the compute path; if they do, diff --git a/src/shinro/codegen/tracing.py b/src/shinro/codegen/tracing.py index 62b62f6..9637f15 100644 --- a/src/shinro/codegen/tracing.py +++ b/src/shinro/codegen/tracing.py @@ -162,6 +162,22 @@ def __neg__(self) -> Tracer: node = self._g.emit("neg", [self.node], self.shape) return Tracer(self._g, self.shape, node) + def __pow__(self, other: Any) -> Tracer: + # The exponent is usually a concrete config scalar (e.g. SMC's alpha), + # which _lift freezes as a 0-d const — the graph stores one specialized + # power. Emitting a `pow` node (rather than numerically baking x**alpha) + # keeps the exponent a compile-time literal while the base stays live. + other = _lift(self._g, other) + out_shape = _broadcast_shape(self.shape, other.shape) + node = self._g.emit("pow", [self.node, other.node], out_shape) + return Tracer(self._g, out_shape, node) + + def __rpow__(self, other: Any) -> Tracer: + other = _lift(self._g, other) + out_shape = _broadcast_shape(other.shape, self.shape) + node = self._g.emit("pow", [other.node, self.node], out_shape) + return Tracer(self._g, out_shape, node) + def __truediv__(self, other: Any) -> Tracer: other = _lift(self._g, other) out_shape = _broadcast_shape(self.shape, other.shape) @@ -179,6 +195,16 @@ def __ne__(self, other: Any) -> Tracer: # type: ignore[override] node = self._g.emit("ne", [self.node, other.node], out_shape) return Tracer(self._g, out_shape, node) + def __lt__(self, other: Any) -> Tracer: # type: ignore[override] + # Ordered comparison recorded as an `lt` node producing 1.0/0.0 flags — + # `ne`'s ordered sibling, the predicate behind threshold guards (e.g. + # "is |c^T g| below eps?"). Like `__ne__`, deliberately returns a + # Tracer: under tracing `<` is data flow, not a Python boolean. + other = _lift(self._g, other) + out_shape = _broadcast_shape(self.shape, other.shape) + node = self._g.emit("lt", [self.node, other.node], out_shape) + return Tracer(self._g, out_shape, node) + @property def T(self) -> Tracer: """Transpose — reverses the shape.""" diff --git a/src/shinro/components.py b/src/shinro/components.py index 99b5870..826dd88 100644 --- a/src/shinro/components.py +++ b/src/shinro/components.py @@ -354,19 +354,32 @@ def post_engine_step(self, engine: "PhysicsEngine") -> None: """ return None - def dynamics(self, state: Any, control: Any) -> Any: + def dynamics(self, state: Any, control: Any, bk: Any | None = None) -> Any: """Continuous-time dynamics :math:`\\dot{x} = f(x, u)`. Override in nonlinear plants to expose the dynamics function for - linearization. Returns None by default (linear plants need not - override this). + linearization *and* for batched rollouts. Returns None by default + (linear plants need not override this). + + The override must accept a single state ``(n_x,)`` **or** a batch + ``(N, n_x)`` and return the derivative with the same rank — a single + state is a batch of one, so one implementation serves both. It must + also route every backend call through ``bk`` (which defaults to + ``self.bk``), because tracing swaps only the traced component's + backend, never the plant's. That is what lets the *same* function back + the eager per-sample rollout, the finite-difference linearization, and + the lowered graph — there is no second, separately-maintained batched + formula to drift. See :mod:`shinro.utils.batching` for the rank + helpers. Args: - state: Current state vector (n_x,). - control: Control input vector (n_u,). + state: Current state vector (n_x,) or batch (N, n_x). + control: Control input (n_u,), batch (N, n_u), or scalar. + bk: Backend to evaluate with. Defaults to ``self.bk``. Returns: - Time derivative of the state (n_x,), or None if not implemented. + Time derivative with the rank of ``state``, or None if not + implemented. """ return None diff --git a/src/shinro/configs/controllers/mppi_inverted_pendulum.toml b/src/shinro/configs/controllers/mppi_inverted_pendulum.toml new file mode 100644 index 0000000..ad4a65e --- /dev/null +++ b/src/shinro/configs/controllers/mppi_inverted_pendulum.toml @@ -0,0 +1,20 @@ +# MPPI for the InvertedPendulum plant — 2-state [theta, theta_dot], 1-D torque +# control [tau]. D_u is inferred from noise_sigma's length, so the bounds and +# costs must all be length 1. +# +# attach_plant wires the rollout to the plant's batch-capable `dynamics`, so +# this config lowers with a *nonlinear* rollout: the graph carries sin/mul +# nodes over the sample batch rather than one linear matmul. Everything else +# (the epsilon port, softmax, the recurrent nominal plan) is unchanged. +type = "MPPI" +name = "inverted_pendulum_mppi" +dt = 0.01 +horizon = 15 +num_samples = 200 +temperature = 0.5 +noise_sigma = [0.5] +state_cost = [20.0, 1.0] +control_cost = [0.1] +u_min = [-0.5] +u_max = [0.5] +seed = 42 diff --git a/src/shinro/controllers/mppi.py b/src/shinro/controllers/mppi.py index 041fa27..9e4678e 100644 --- a/src/shinro/controllers/mppi.py +++ b/src/shinro/controllers/mppi.py @@ -25,10 +25,14 @@ :class:`BatchedDynamicsAdapter`, which vectorizes the plant's single-state model over the sample batch. -The Gaussian sampling and softmax weighting are numpy-based (no RNG -abstraction in ``ArrayBackend``). They run on CPU and are converted to the -backend via ``bk.from_numpy``, following the same bridge pattern MPC uses for -OSQP. +Sampling is the one step of a tick with no dataflow representation, so it +stays on the host: perturbations are drawn with numpy (``ArrayBackend`` has no +RNG abstraction) and bridged to the backend via ``bk.from_numpy``. Everything +after sampling — the batched rollout, the softmax weighting, the +nominal-sequence update — runs through ``self.bk``, which is what makes the +controller lowerable: a traced ``compute`` receives the already-drawn +perturbations through a free graph input port (``epsilon``) instead of +sampling, and the compiled kernel computes the rest. Usage: controller = MPPIController( @@ -55,6 +59,30 @@ from shinro.utils.array_backend import ArrayBackend, NumpyBackend +def _as_float(value: Any, field: str) -> float: + """Coerce a config scalar to float, naming the field when it fails. + + Mirrors :func:`shinro.controllers.smc._as_float`: config values arrive from + TOML (or hand-written dicts), so a typo is a user error worth naming — a + bare ``float("abc")`` reports only that a conversion failed, not which + field or config was wrong. + + Args: + value: The raw config value. + field: The config field name, for the error message. + + Returns: + The value as a ``float``. + + Raises: + ValueError: If the value is not numeric. + """ + try: + return float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"MPPIConfig.{field} must be a number, got {value!r}") from exc + + @dataclass(frozen=True) class MPPIConfig: """Strict TOML schema for :class:`MPPIController`. @@ -144,8 +172,8 @@ def __init__( self.N = num_samples self.K = horizon - self.dt = float(dt) - self.lam = float(temperature) + self.dt = _as_float(dt, "dt") + self.lam = _as_float(temperature, "temperature") self.seed = seed self._rng = np.random.default_rng(seed) @@ -164,16 +192,38 @@ def __init__( self._adapter = None self._Q = None self._R = None - self._x_ref = None + # Tracking reference for the attach_plant cost. Held in a 1-element + # list rather than as an array-like attr: the tracer's attr-diff + # promotes every reassigned ndarray/Tracer attr to a recurrent state + # port, and this is a per-call input, not state. + self._x_ref_holder = [None] + + def host_input_shapes(self) -> dict[str, tuple[int, ...]]: + """Free C-ABI input ports the host fills each tick. + + Sampling is host-side (see the module docstring), so in a lowered graph + ``epsilon`` — the ``(N, K*D_u)`` sampled perturbations, sample-major — + is a free input port rather than something the kernel generates. The + compile pipeline (:func:`shinro.codegen.build.build_composed_graph`) + reads this to declare the port and pass it through + ``compose(host_inputs=...)``; the eager path ignores it, because + :meth:`compute` samples internally. + + Returns: + Maps the port name to its shape. + """ + return {"epsilon": (self.N, self.K * self.D_u)} def attach_plant(self, plant, Q: Any | None = None, R: Any | None = None): """Wire a plant into the controller via a batched dynamics adapter. Builds a :class:`BatchedDynamicsAdapter` from the plant and sets ``dynamics_fn`` / ``cost_fn`` from it. The cost uses the quadratic - stage cost :math:`x^T Q x + u^T R u`, optionally tracking a reference - passed to :func:`compute`. ``Q`` and ``R`` may be diagonal ``(D,)`` or - full ``(D, D)`` matrices; they default to identity. + stage cost :math:`x^T Q x + u^T R u`, optionally tracking the reference + passed to :func:`compute` — that reference is threaded through + ``self._x_ref_holder``, a plain list, so the tracer does not mistake it + for recurrent state. ``Q`` and ``R`` may be diagonal ``(D,)`` or full + ``(D, D)`` matrices; they default to identity. Args: plant: A :class:`Plant` exposing ``get_model()`` (and optionally @@ -199,21 +249,27 @@ def attach_plant(self, plant, Q: Any | None = None, R: Any | None = None): self.D_x = adapter.state_dim self.D_u = adapter.control_dim - self.dynamics_fn = adapter.dynamics_fn - self.cost_fn = lambda x, u: adapter.cost_fn(x, u, self._Q, self._R, x_ref=self._x_ref) + # Route the controller's *current* backend into the batched path: a + # trace swaps self.bk to a TraceBackend, while the adapter holds the + # plant's (concrete) backend. Resolving self.bk at call time is what + # lets a plant's batched_dynamics emit sin/mul/slice nodes instead of + # evaluating eagerly against a concrete array. + self.dynamics_fn = lambda x, u, dt: adapter.dynamics_fn(x, u, dt, bk=self.bk) + self.cost_fn = lambda x, u: adapter.cost_fn(x, u, self._Q, self._R, x_ref=self._x_ref_holder[0]) - def compute(self, current_state, target_state: Any | None = None): + def compute(self, current_state, target_state: Any | None = None, epsilon: Any | None = None): """Compute the MPPI control action for a given initial state. - Samples :math:`N` Gaussian perturbation sequences, rolls out the - dynamics over the horizon, computes the softmax-weighted update, and - returns the first action of the updated nominal sequence (clipped to - bounds if configured). + Draws :math:`N` Gaussian perturbation sequences (host-side — see the + module docstring), rolls out the dynamics over the horizon, computes + the softmax-weighted update, and returns the first action of the + updated nominal sequence (clipped to bounds if configured). - The rollout loop runs on ``self.bk``, so with a :class:`TorchBackend` - the batched dynamics/cost operations execute as torch tensor ops. The - Gaussian sampling and softmax weighting run in numpy and are bridged - to the backend. + Everything after sampling runs through ``self.bk`` on batched tensors, + so it evaluates identically on numpy, torch, and the tracing backend. + Every reduction is written as a matmul identity — a sum over an axis is + a contraction with a ones vector or with the softmax weights — except + the softmax shift ``beta = min(costs)``, which is the ``min`` graph op. Args: current_state: Initial state vector (D_x,). Accepts the backend's @@ -221,6 +277,12 @@ def compute(self, current_state, target_state: Any | None = None): target_state: Optional reference state (D_x,) to track. When given, the cost penalizes deviation ``(x - target_state)``; otherwise the controller regulates to the origin. + epsilon: Optional pre-drawn perturbations, shape ``(N, K*D_u)``, + sample-major (row ``i``, column ``k*D_u + d``). This is the + trace hook: when lowering, the perturbations become a free + graph input port the host fills each tick, so the traced call + never samples. ``None`` (the eager path) draws them from + ``self._rng``. Returns: First control action (D_u,) in the backend's native type. @@ -238,58 +300,109 @@ def compute(self, current_state, target_state: Any | None = None): dynamics_fn = self.dynamics_fn cost_fn = self.cost_fn - x0_np = self.bk.to_numpy(current_state) - self._x_ref = self.bk.from_numpy(self.bk.to_numpy(target_state)) if target_state is not None else None - - epsilon = self._rng.normal(loc=0.0, scale=self.noise_sigma, size=(self.N, self.K, self.D_u)) - self._last_epsilon = epsilon - - v = np.expand_dims(self.u, axis=0) + epsilon - if self.u_min is not None or self.u_max is not None: - v = np.clip(v, self.u_min, self.u_max) - - # Backend-native rollout. The dynamics/cost callables operate on the - # backend's tensors; the sampled perturbations and nominal sequence - # live in numpy and are bridged once before the loop. Each from_numpy - # is a CPU->GPU transfer on a torch backend, so converting the - # loop-invariant arrays up front removes 2K transfers per compute(). - x_current = self.bk.from_numpy(np.tile(x0_np, (self.N, 1))) - u_plan = self.bk.from_numpy(v) - eps_b = self.bk.from_numpy(epsilon) - u_nom_b = self.bk.from_numpy(self.u) - sigma2_b = self.bk.from_numpy(self.noise_sigma**2) - costs = self.bk.zeros(self.N) + # Sampling is host-side by design: the one step of the tick with no + # dataflow representation. The traced path receives epsilon as a graph + # input port, so this branch is a trace-time constant and the traced + # call never touches the RNG. + if epsilon is None: + eps_sample = self._rng.normal(loc=0.0, scale=self.noise_sigma, size=(self.N, self.K, self.D_u)) + self._last_epsilon = eps_sample + eps_in = self.bk.from_numpy(eps_sample.reshape(self.N, self.K * self.D_u)) + else: + # Host-supplied perturbations (the trace path, or a caller wanting + # reproducible noise). Bridged like the sampled path so a numpy + # array works on a torch backend; a no-op under tracing, where + # epsilon is already the graph input Tracer. + eps_in = self.bk.from_numpy(epsilon) + + N, K, Du = self.N, self.K, self.D_u lam = self.lam - for k in range(self.K): - u_k = u_plan[:, k, :] - costs = costs + cost_fn(x_current, u_k) - inv_var_weighted_u = u_nom_b[k] / sigma2_b - control_penalty = lam * self.bk.sum(inv_var_weighted_u * eps_b[:, k, :], axis=1) + # Batch the initial state to (N, D_x) with a same-rank broadcast; the + # tracer rejects numpy's rank-differing tile. + d_x = current_state.shape[0] + x_current = self.bk.zeros((N, d_x)) + self.bk.reshape(current_state, (1, d_x)) + + # Tracking reference for the attach_plant cost closure (see __init__), + # stored as a (1, D_x) row: the closure subtracts it from the batched + # states, and (N, D_x) - (1, D_x) stays same-rank for the tracer (numpy + # would broadcast a (D_x,) there; the tracer rejects the rank gap). + self._x_ref_holder[0] = self.bk.reshape(target_state, (1, d_x)) if target_state is not None else None + + # eps_in is (N, K*D_u) sample-major, so its transpose is (K*D_u, N) and + # step k's perturbations are the row block [k*D_u:(k+1)*D_u] transposed + # back to (N, D_u). + eps_t = eps_in.T + ones_du = self.bk.array(np.ones(Du)) + ones_n_col = self.bk.array(np.ones((N, 1))) + sigma2 = self.bk.reshape(self.bk.array(self.noise_sigma**2), (1, Du)) + costs = self.bk.zeros(N) + # Bridge the nominal sequence to the backend once. It is numpy-hosted + # (see the assignment at the end), and `numpy + torch` raises — only + # `torch + numpy` works — so every use goes through this backend-native + # copy. Under tracing this is a no-op passthrough and u_nominal is the + # recurrent state input Tracer. + u_nominal = self.bk.from_numpy(self.u) + + for k in range(K): + eps_k = self.bk.slice_(eps_t, k * Du, (k + 1) * Du).T + u_k = self.bk.slice_(u_nominal, k, k + 1) # nominal row -> (1, D_u) + v_k = eps_k + u_k + if self.u_min is not None or self.u_max is not None: + v_k = self.bk.clip(v_k, self.u_min, self.u_max) + costs = costs + cost_fn(x_current, v_k) + # lam * sum_d (u_k,d / sigma_d^2) eps_i,k,d — a row sum, i.e. a + # contraction with a ones vector. + control_penalty = lam * ((u_k / sigma2 * eps_k) @ ones_du) costs = costs + control_penalty - x_current = dynamics_fn(x_current, u_k, self.dt) - - costs = costs + cost_fn(x_current, self.bk.zeros((self.N, self.D_u))) - costs_np = self.bk.to_numpy(costs) - self._last_costs = costs_np.copy() - - beta = np.min(costs_np) - softmax_w = np.exp(-(costs_np - beta) / lam) - softmax_w /= np.sum(softmax_w) - - weighted_eps = np.sum(softmax_w[:, np.newaxis, np.newaxis] * epsilon, axis=0) + x_current = dynamics_fn(x_current, v_k, self.dt) - self.u += weighted_eps - - u_0 = self.u[0].copy() - if self.K > 1: - self.u[:-1] = self.u[1:] - self.u[-1] = self.u[-2] + costs = costs + cost_fn(x_current, self.bk.zeros((N, Du))) + # Diagnostics: the traced backend publishes `costs` as an auxiliary + # graph output port (SMC's `healthy` precedent); eager backends keep the + # concrete array. Assigning a Tracer here would be mis-detected as + # recurrent state. + self.bk.emit_named_output("costs", costs) + costs_np = self.bk.to_numpy(costs) + if isinstance(costs_np, np.ndarray): + self._last_costs = costs_np.copy() + + # Softmax weights. beta is the shift that keeps exp from overflowing. + # The normalizer is a contraction with a ones vector, not a 1-D @ 1-D + # product: (1,N) @ (N,1) with the weights on *both* sides would sum the + # squares. The (1,1) result ravels to (1,), which broadcasts against (N,). + beta = self.bk.min(costs) + w = self.bk.exp(-(costs - beta) / lam) + w_sum = self.bk.ravel(self.bk.reshape(w, (1, N)) @ ones_n_col) + w = w / w_sum + + # Weighted average of the perturbations: sum_i w_i eps_i is the matmul + # (1,N) @ (N, K*D_u). to_numpy keeps the nominal sequence numpy-hosted + # on the eager backends; under tracing it is a no-op passthrough, so + # self.u stays a Tracer and is emitted as the recurrent state output. + w_row = self.bk.reshape(w, (1, N)) + weighted_eps = self.bk.reshape(w_row @ eps_in, (K, Du)) + u_updated = u_nominal + weighted_eps + + # The returned action is the first element *before* the receding-horizon + # shift (self.u[:-1] = self.u[1:]; self.u[-1] = self.u[-2], i.e. rows + # 1..K-1 followed by row K-1 again). The tracer has no in-place slice + # assignment, so the shift is K row slices + a stack. + u_0 = self.bk.slice_(u_updated, 0, 1) + if K > 1: + rows = [self.bk.ravel(self.bk.slice_(u_updated, j, j + 1)) for j in range(1, K)] + rows.append(self.bk.ravel(self.bk.slice_(u_updated, K - 1, K))) + shifted = self.bk.stack(rows) + else: + shifted = u_updated + self.u = self.bk.to_numpy(shifted) + + u_0 = self.bk.ravel(u_0) if self.u_min is not None or self.u_max is not None: - u_0 = np.clip(u_0, self.u_min, self.u_max) + u_0 = self.bk.clip(u_0, self.u_min, self.u_max) - return self.bk.from_numpy(u_0) + return u_0 def reset(self): """Reset the controller to its initial state. diff --git a/src/shinro/controllers/smc.py b/src/shinro/controllers/smc.py index 392b9e6..6e5b03e 100644 --- a/src/shinro/controllers/smc.py +++ b/src/shinro/controllers/smc.py @@ -9,6 +9,22 @@ where :math:`s = c^T x` is the sliding surface. Supports multiple boundary-layer smoothers (sat, tanh, sigmoid) to suppress chattering. +Scope — a single sliding surface, for every plant shape: + ``c`` is one row, so ``s = c^T x`` is a scalar and the law controls one + error combination. This covers the whole ``(n_x, n_u)`` rectangle: + ``n_u == 1`` divides directly (:meth:`SlidingModeController._scalar_control`), + ``n_u > 1`` takes the minimum-norm closed form of the underdetermined + system ``(c^T g) u = num`` (:meth:`SlidingModeController._vector_control`), + and both behave the same when evaluated and when lowered to a compiled + kernel. **Multi-surface SMC is deliberately out of scope** — a surface + *matrix* ``C`` (``m > 1`` surfaces solved jointly) is not implemented. + Field practice is single-surface, or decentralized (one instance per axis, + which this class already supports by instantiation), and for coupled MIMO + plants the min-norm solve would not respect actuator limits anyway (that is + an allocation problem, not a control-law one). The gate for revisiting it + is the off-diagonal/diagonal ratio of ``C·g(x)`` at the operating point; + see the 2026-09-14 lab note for the full decision record. + Usage: # In configs/controllers/smc.toml: # type = "SMC" @@ -21,6 +37,7 @@ """ from dataclasses import dataclass +from typing import Any import numpy as np @@ -29,6 +46,29 @@ from shinro.utils.array_backend import ArrayBackend, NumpyBackend +def _as_float(value: Any, field: str) -> float: + """Coerce a config scalar to float, naming the field when it fails. + + Config values arrive from TOML (or hand-written dicts), so a typo is a + user error worth naming: a bare ``float("abc")`` reports only "could not + convert string to float", not which field or config was wrong. + + Args: + value: The raw config value. + field: The config field name, for the error message. + + Returns: + The value as a ``float``. + + Raises: + ValueError: If the value is not numeric. + """ + try: + return float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"SMCConfig.{field} must be a number, got {value!r}") from exc + + @dataclass(frozen=True) class SMCConfig: """Strict TOML schema for :class:`SlidingModeController`. @@ -43,6 +83,7 @@ class SMCConfig: k2: float = 0.0 smoother: str = "sat" alpha: float = 0.0 + controllability_eps: float = 1e-12 dt: float | None = None name: str = "smc" @@ -55,6 +96,11 @@ class SlidingModeController(Controller): optional boundary-layer smoothing. The sliding surface coefficients ``c`` must form a Hurwitz polynomial. + One surface per instance: ``c`` is a single surface vector, so the class is + instantiated once per controlled axis (a bank of instances is the + multi-axis pattern). Multi-surface SMC — a surface matrix solved jointly — + is intentionally not implemented; see the module docstring's "Scope" note. + Args: c: Sliding surface coefficients (n,). The polynomial ``c[0] + c[1] p + ... + c[n-1] p^{n-1}`` must be Hurwitz. @@ -66,6 +112,19 @@ class SlidingModeController(Controller): ``"tanh"``, or ``"sigmoid"``. alpha: Fractional power exponent for the switching term :math:`|s|^\\alpha`. 0 gives sign-only; 1 gives linear. + controllability_eps: Near-zero threshold on the controllability + denominator: :math:`|c^T g(x)|` for a scalar input, and its + multi-input generalization :math:`\\|c^T g(x)\\|` (the two agree + when ``n_u == 1``). **A deployment design parameter, not a + numerical constant** — set it above the smallest value the plant + can legitimately produce (the law amplifies :math:`1/c^T g`, so + command saturation and chattering arrive long before the + arithmetic floor). The 1e-12 default only protects the division + itself, for plants whose ``g`` is well-conditioned everywhere. + Live numpy calls raise ``RuntimeError`` below it; the lowered + (compiled) graph instead emits a fail-safe zero command and + reports a ``healthy`` flag, since a straight-line kernel cannot + raise. backend: Array backend. Defaults to NumpyBackend. """ @@ -77,14 +136,16 @@ def __init__( k2: float = 0.0, smoother: str = "sat", alpha: float = 0.0, + controllability_eps: float = 1e-12, backend: ArrayBackend | None = None, ): self.bk = backend or NumpyBackend() self.c = self.bk.array(c).flatten() - self.k1 = float(k1) - self.k2 = float(k2) - self.phi = float(phi) - self.alpha = float(alpha) + self.k1 = _as_float(k1, "k1") + self.k2 = _as_float(k2, "k2") + self.phi = _as_float(phi, "phi") + self.alpha = _as_float(alpha, "alpha") + self.controllability_eps = _as_float(controllability_eps, "controllability_eps") SMOOTHERS = self._dict_boundaries() @@ -139,7 +200,18 @@ def compute(self, x, f_x, g_x): Evaluates :math:`u = (c^T g)^{-1} ( -c^T f - k_1 |s|^\\alpha \\, \\text{smooth}(s) - k_2 s )`. For scalar input (``c^T g`` is scalar), uses direct division. For - vector input, solves the least-squares problem. + vector input the equivalent-control equation ``(c^T g) u = num`` is + underdetermined (one surface, ``n_u`` unknowns), so the law takes its + minimum-norm solution ``u = cg^T (cg cg^T)^{-1} num`` — the same + point ``np.linalg.lstsq`` returned, but written as matmul/transpose/ + divide so it traces and lowers. + + The computation is written so it traces as well as it evaluates: the + ``u``/``n_u`` branch keys off ``g_x``'s *shape* (known at trace time), + the dot products use column-vector form (the tracer has no 1D @ 1D + contraction), and the near-zero ``c^T g`` guard is an exception only + when ``g_x`` is concrete — a graph cannot raise, so the traced path + emits a fail-safe instead (see :meth:`_scalar_control`). Args: x: Current state vector (n,). @@ -150,15 +222,24 @@ def compute(self, x, f_x, g_x): Control input vector (n_u,). Raises: - RuntimeError: If :math:`c^T g(x)` is near-zero for scalar input. + RuntimeError: If :math:`c^T g(x)` is near-zero and ``g_x`` is + concrete (eager backends only; the traced path cannot raise). """ - x = self.bk.array(x).flatten() - f_x = self.bk.array(f_x).flatten() + x = self.bk.ravel(self.bk.array(x)) + f_x = self.bk.ravel(self.bk.array(f_x)) g_x = self.bk.array(g_x) - - s = self.c @ x - cf = self.c @ f_x - cg = self.c @ g_x + if len(g_x.shape) == 1: + # 1-D g_x is a single control column; normalize to (n, 1). + g_x = self.bk.reshape(g_x, (g_x.shape[0], 1)) + + n = self.n + # Column-vector contractions: (n,) @ (n, 1) -> (1,) and + # (n,) @ (n, n_u) -> (n_u,). numpy would collapse 1D @ 1D to a 0-d + # scalar, which the tracer rejects (no 0-d propagation) — same numbers, + # a uniform (1,)/(n_u,) shape on every backend. + s = self.c @ self.bk.reshape(x, (n, 1)) + cf = self.c @ self.bk.reshape(f_x, (n, 1)) + cg = self.bk.ravel(self.c @ g_x) if self.phi > 0: smooth_s = self._smoother(s) @@ -166,21 +247,106 @@ def compute(self, x, f_x, g_x): smooth_s = self.bk.sign(s) s_dot_desired = -self.k1 * self.bk.abs(s) ** self.alpha * smooth_s - self.k2 * s + num = s_dot_desired - cf + + if len(g_x.shape) >= 2 and g_x.shape[-1] > 1: + return self._vector_control(num, cg) + return self._scalar_control(num, cg) + + def _scalar_control(self, num, cg): + """Single-input law: ``u = (s_dot_desired - c^T f) / (c^T g)``. + + The guard behaves differently per backend by design: + + - **Concrete (numpy/torch)**: ``c^T g`` below ``controllability_eps`` + raises ``RuntimeError`` — the developer-facing signal. + - **Traced**: a graph cannot raise (the compiled VM is straight-line; + a Zig panic crossing the C ABI aborts the host process). Instead the + guard becomes data: ``u`` is forced to the fail-safe zero command + and a ``healthy`` flag is published as an auxiliary output port + (:meth:`ArrayBackend.emit_named_output`), so the host can run its own + fault policy inside the same tick. Zero-command is the floor, not a + safety guarantee — for an open-loop-unstable plant the host policy is + the real fail-safe. - cg_flat = self.bk.ravel(cg) - cg_size = self.bk.to_numpy(cg_flat).size - if cg_size == 1: - cg_val = float(self.bk.to_numpy(cg_flat)[0]) - if abs(cg_val) < 1e-12: + Args: + num: ``s_dot_desired - c^T f`` (1,). + cg: ``c^T g`` (1,). + + Returns: + Control input (1,). + + Raises: + RuntimeError: If ``|c^T g| < controllability_eps`` and ``cg`` is + concrete. + """ + concrete = self.bk.to_numpy(cg) + if isinstance(concrete, np.ndarray): + cg_val = np.asarray(concrete).reshape(-1)[0] + if abs(cg_val) < self.controllability_eps: raise RuntimeError("c^T g(x) is near-zero — loss of controllability") - u = self.bk.array([(s_dot_desired - cf) / cg_val]) - else: - cg_np = self.bk.to_numpy(cg_flat).reshape(1, -1) - rhs_np = np.array([float(self.bk.to_numpy(s_dot_desired - cf))]) - u_np, _, _, _ = np.linalg.lstsq(cg_np, rhs_np, rcond=None) - u = self.bk.from_numpy(u_np.flatten()) + return num / cg + + # Traced: the guard is data, not an exception (see the class docstring + # above). u_raw may compute inf/nan here; `where` discards it. + u_raw = num / cg + fault = self.bk.abs(cg) < self.controllability_eps + self.bk.emit_named_output("healthy", 1.0 - fault) + return self.bk.where(fault, self.bk.zeros_like(u_raw), u_raw) + + def _vector_control(self, num, cg): + """Multi-input law: minimum-norm solution of ``(c^T g) u = num``. + + ``c^T g`` is a ``1 x n_u`` row, so the equation is underdetermined and + ``(c^T g)^{-1}`` does not exist. Writing ``A = c^T g``, the + minimum-norm exact solution is the pseudo-inverse ``A^+ b``, which for + a single full-rank row collapses to ``A^T (A A^T)^{-1} b`` — the same + answer ``np.linalg.lstsq`` gives, but built from matmul/transpose/ + divide so both the eager backends and the traced/lowered graph share + one implementation. + + The guard mirrors :meth:`_scalar_control` (see that docstring for why + the concrete and traced backends diverge): it tests ``‖c^T g‖`` + rather than its square, so ``controllability_eps`` means the same + magnitude here as ``|c^T g|`` does on the scalar branch (``n_u == 1`` + makes them identical). + + Args: + num: ``s_dot_desired - c^T f`` (1,). + cg: ``c^T g`` (n_u,). + + Returns: + Control input (n_u,). + + Raises: + RuntimeError: If ``‖c^T g‖ < controllability_eps`` and ``cg`` + is concrete. + """ + n_u = cg.shape[0] + # Make the row explicit so A^T (A A^T)^-1 is expressible with matmul + # and transpose alone (the tracer has no 1-D pseudo-inverse op). + cg_row = self.bk.reshape(cg, (1, n_u)) + denom = cg_row @ cg_row.T # (1,1) — ||cg||^2, the Gram scalar + norm = denom**0.5 # (1,1) — ||cg||; equals |c^T g| when n_u == 1 + num_col = self.bk.reshape(num, (1, 1)) # (1,) -> (1,1), rank-aligned + + concrete = self.bk.to_numpy(norm) + if isinstance(concrete, np.ndarray): + # Guard before dividing: on the invalid input the division would + # only produce inf/nan (and a numpy divide warning) before the raise. + norm_val = np.asarray(concrete).reshape(-1)[0] + if norm_val < self.controllability_eps: + raise RuntimeError("c^T g(x) is near-zero — loss of controllability") + return self.bk.ravel(cg_row.T @ (num_col / denom)) - return u + # Traced: the guard is data, not an exception (see _scalar_control). + # Apply it to the (1,1) scalar factor rather than to u: zeroing scale + # zeroes u = cg^T scale, and `where`'s operands stay the same shape. + scale = num_col / denom # (1,1) — the A^T b factor; may be inf/nan + fault = norm < self.controllability_eps + self.bk.emit_named_output("healthy", 1.0 - self.bk.ravel(fault)) + safe = self.bk.where(fault, self.bk.zeros_like(scale), scale) + return self.bk.ravel(cg_row.T @ safe) def reset(self): """No internal state to reset for SMC.""" @@ -199,6 +365,9 @@ def from_config(cls, config, backend: ArrayBackend | None = None): smoother: Smoothing function — ``"sat"``, ``"tanh"``, or ``"sigmoid"`` (default ``"sat"``). alpha: Fractional power exponent (default 0.0). + controllability_eps: Near-zero :math:`|c^T g|` threshold, + plant-scaled (default 1e-12 = arithmetic-only guard); see + :class:`SlidingModeController` for how to choose it. Args: config: TOML config dict or SMCConfig. @@ -216,5 +385,6 @@ def from_config(cls, config, backend: ArrayBackend | None = None): k2=cfg.k2, smoother=cfg.smoother, alpha=cfg.alpha, + controllability_eps=cfg.controllability_eps, backend=bk, ) diff --git a/src/shinro/plants/cartpole.py b/src/shinro/plants/cartpole.py index 1d03785..a7318fa 100644 --- a/src/shinro/plants/cartpole.py +++ b/src/shinro/plants/cartpole.py @@ -3,6 +3,7 @@ from shinro.components import PhysicsEngine, Plant from shinro.factories.registry import register_plant, register_plant_detector from shinro.utils.array_backend import ArrayBackend, NumpyBackend +from shinro.utils.batching import as_batch, as_vector, column, control_batch from shinro.utils.config_spec import strip_runtime_keys from shinro.utils.linearization import discretize_euler, linearize_plant @@ -163,7 +164,7 @@ def get_model(self, x0=None, u0=None, eps=1e-6): A_c, B_c = linearize_plant(self, x0, u0, eps=eps) return discretize_euler(A_c, B_c, self.dt, backend=self.bk) - def _compute_accels(self, x, theta, x_dot, theta_dot, F): + def _compute_accels(self, x, theta, x_dot, theta_dot, F, bk=None): """Compute the accelerations from the equations of motion. Solves the coupled 2x2 system for :math:`\\ddot{x}` and @@ -175,32 +176,43 @@ def _compute_accels(self, x, theta, x_dot, theta_dot, F): x_dot: Cart velocity (m/s). theta_dot: Pole angular velocity (rad/s). F: Horizontal force on cart (N). + bk: Backend to evaluate with. Defaults to the plant's backend. Returns: Tuple of (x_ddot, theta_ddot). """ + bk = self.bk if bk is None else bk M, m, pole_len, g, _ = self.M, self.m, self.l, self.g, self.b - sin_theta = self.bk.sin(theta) - cos_theta = self.bk.cos(theta) + sin_theta = bk.sin(theta) + cos_theta = bk.cos(theta) denom = pole_len - m * pole_len * cos_theta**2 / (M + m) theta_ddot = (g * sin_theta - cos_theta * (F + m * pole_len * theta_dot**2 * sin_theta) / (M + m)) / denom x_ddot = (F + m * pole_len * (theta_dot**2 * sin_theta - theta_ddot * cos_theta)) / (M + m) return x_ddot, theta_ddot - def dynamics(self, state, control): + def dynamics(self, state, control, bk=None): """Continuous-time dynamics :math:`\\dot{x} = f(x, u)`. Args: - state: State vector (4,) — [x, x_dot, theta, theta_dot]. - control: Control vector (1,) or scalar — [F]. + state: State vector (4,) — [x, x_dot, theta, theta_dot] — or a + batch (N, 4). + control: Control (1,), batch (N, 1), or scalar — [F]. + bk: Backend to evaluate with. Defaults to the plant's backend. Returns: - Time derivative of the state (4,) — [x_dot, x_ddot, theta_dot, theta_ddot]. + Time derivative with the rank of ``state`` — + [x_dot, x_ddot, theta_dot, theta_ddot]. """ - x, x_dot, theta, theta_dot = state[0], state[1], state[2], state[3] - F = control[0] if hasattr(control, '__len__') else control - x_ddot, theta_ddot = self._compute_accels(x, theta, x_dot, theta_dot, F) - return self.bk.stack([x_dot, x_ddot, theta_dot, theta_ddot]) + bk = self.bk if bk is None else bk + x, single = as_batch(bk, state) + u = control_batch(bk, control, 1) + x_pos = column(bk, x, 0) + x_dot = column(bk, x, 1) + theta = column(bk, x, 2) + theta_dot = column(bk, x, 3) + x_ddot, theta_ddot = self._compute_accels(x_pos, theta, x_dot, theta_dot, column(bk, u, 0), bk=bk) + f = bk.stack([bk.ravel(x_dot), bk.ravel(x_ddot), bk.ravel(theta_dot), bk.ravel(theta_ddot)]).T + return as_vector(bk, f, single) def step(self, u): """Execute one control step. @@ -221,13 +233,11 @@ def step(self, u): self.state = self.get_state() return self.state - x, x_dot, theta, theta_dot = self.state[0], self.state[1], self.state[2], self.state[3] - F = u[0] if hasattr(u, '__len__') else u - x_ddot, theta_ddot = self._compute_accels(x, theta, x_dot, theta_dot, F) - theta_dot_new = theta_dot + theta_ddot * self.dt - x_dot_new = x_dot + x_ddot * self.dt - theta_new = theta + theta_dot_new * self.dt - x_new = x + x_dot_new * self.dt + f = self.dynamics(self.state, u) + theta_dot_new = self.state[3] + f[3] * self.dt + x_dot_new = self.state[1] + f[1] * self.dt + theta_new = self.state[2] + theta_dot_new * self.dt + x_new = self.state[0] + x_dot_new * self.dt self.state = self.bk.array([x_new, x_dot_new, theta_new, theta_dot_new]) if self.track_limits is not None: self.state = self.bk.array([ diff --git a/src/shinro/plants/double_pendulum.py b/src/shinro/plants/double_pendulum.py index 4ab20d4..5ae5dea 100644 --- a/src/shinro/plants/double_pendulum.py +++ b/src/shinro/plants/double_pendulum.py @@ -3,6 +3,7 @@ from shinro.components import PhysicsEngine, Plant from shinro.factories.registry import register_plant, register_plant_detector from shinro.utils.array_backend import ArrayBackend, NumpyBackend +from shinro.utils.batching import as_batch, as_vector, column, control_batch from shinro.utils.config_spec import BoundsConfig, strict_from_dict, strip_runtime_keys from shinro.utils.linearization import discretize_euler, linearize_plant @@ -115,81 +116,81 @@ def physics_engine(self, engine: PhysicsEngine | None): self.bk = NumpyBackend() self.state = self.bk.zeros(4) - def _make_mass_matrix(self, diff_theta): - """Build the 2x2 mass matrix :math:`M(\\theta)`. + def dynamics(self, state, control, bk=None): + """Continuous-time dynamics :math:`\\dot{x} = f(x, u)`, batch-capable. - Args: - diff_theta: Angle difference :math:`\\theta_1 - \\theta_2`. - - Returns: - Mass matrix M (2, 2). - """ - M = self.bk.zeros((2, 2)) - M[0, 0] = (self.m1 + self.m2) * self.l1**2 - M[0, 1] = self.m2 * self.l1 * self.l2 * self.bk.cos(diff_theta) - M[1, 0] = M[0, 1] - M[1, 1] = self.m2 * self.l2**2 - return M - - def _make_coriolis_matrix(self, diff_theta, angular_velocities): - """Build the 2x2 Coriolis matrix :math:`C(\\theta, \\dot{\\theta})`. - - Args: - diff_theta: Angle difference :math:`\\theta_1 - \\theta_2`. - angular_velocities: Vector (2,) — [omega_1, omega_2]. - - Returns: - Coriolis matrix C (2, 2). - """ - C = self.bk.zeros((2, 2)) - w_1, w_2 = angular_velocities[0], angular_velocities[1] - C[0, 1] = self.m2 * self.l1 * self.l2 * self.bk.sin(diff_theta) * w_2 - C[1, 0] = -self.m2 * self.l2 * self.l1 * self.bk.sin(diff_theta) * w_1 - return C + State ordering is :math:`[\\theta_1, \\theta_2, \\omega_1, \\omega_2]` + and control is :math:`[\\tau_1, \\tau_2]`. The angular acceleration + solves the manipulator equation + :math:`\\ddot{\\theta} = M^{-1}(\\tau - C\\dot{\\theta} - G)`, with - def _make_gravity_vector(self, theta_angles): - """Build the gravity vector :math:`G(\\theta)`. + .. math:: - Args: - theta_angles: Vector (2,) — [theta_1, theta_2]. + M = \\begin{bmatrix} (m_1 + m_2) l_1^2 & m_2 l_1 l_2 \\cos\\Delta \\ + m_2 l_1 l_2 \\cos\\Delta & m_2 l_2^2 \\end{bmatrix}, + \\quad \\Delta = \\theta_1 - \\theta_2, - Returns: - Gravity vector G (2,). - """ - theta_1, theta_2 = theta_angles[0], theta_angles[1] - G = self.bk.zeros(2) - G[0] = (self.m1 + self.m2) * self.g * self.l1 * self.bk.sin(theta_1) - G[1] = self.m2 * self.g * self.l2 * self.bk.sin(theta_2) - return G + :math:`C\\dot{\\theta} = [m_2 l_1 l_2 \\sin\\Delta\\,\\omega_2^2, + -m_2 l_1 l_2 \\sin\\Delta\\,\\omega_1^2]`, and + :math:`G = [(m_1 + m_2) g l_1 \\sin\\theta_1, + m_2 g l_2 \\sin\\theta_2]`. - def dynamics(self, state, control): - """Continuous-time dynamics :math:`\\dot{x} = f(x, u)`. - - State ordering is :math:`[\\theta_1, \\theta_2, \\omega_1, \\omega_2]` - and control is :math:`[\\tau_1, \\tau_2]`. The angular acceleration is - solved from the manipulator equation - :math:`\\ddot{\\theta} = M^{-1}(\\tau - C\\dot{\\theta} - G)`. + The 2x2 solve is written in closed form (Cramer's rule) rather than + ``bk.solve``: on the batched path ``M`` would be rank-3 (``(N, 2, 2)``), + which the graph backend — strictly 2-D — does not represent. The + entries are therefore ``(N, 1)`` columns, and the determinant + :math:`m_2 l_1^2 l_2^2 (m_1 + m_2 \\sin^2\\Delta)` is strictly + positive for positive masses, so the division is always safe. Args: - state: State vector (4,) — [theta_1, theta_2, omega_1, omega_2]. - control: Control vector (2,) or scalar — [tau_1, tau_2]. + state: State (4,) — [theta_1, theta_2, omega_1, omega_2] — or a + batch (N, 4). + control: Control (2,), batch (N, 2), or scalar — [tau_1, tau_2]. + bk: Backend to evaluate with. Defaults to the plant's backend. Returns: - Time derivative of the state (4,) — + Time derivative with the rank of ``state`` — [omega_1, omega_2, theta_1_ddot, theta_2_ddot]. """ - theta_1, theta_2, omega_1, omega_2 = state[0], state[1], state[2], state[3] - diff_theta = theta_1 - theta_2 - omega = self.bk.array([omega_1, omega_2]) - M = self._make_mass_matrix(diff_theta) - C = self._make_coriolis_matrix(diff_theta, omega) - G = self._make_gravity_vector(self.bk.array([theta_1, theta_2])) - - tau = control if hasattr(control, '__len__') else self.bk.array([control, 0.0]) - b = tau - C @ omega - G - thetaddot = self.bk.solve(M, b) - - return self.bk.stack([omega_1, omega_2, thetaddot[0], thetaddot[1]]) + bk = self.bk if bk is None else bk + x, single = as_batch(bk, state) + u = control_batch(bk, control, 2) + theta_1 = column(bk, x, 0) + theta_2 = column(bk, x, 1) + omega_1 = column(bk, x, 2) + omega_2 = column(bk, x, 3) + tau_1 = column(bk, u, 0) + tau_2 = column(bk, u, 1) + + diff = theta_1 - theta_2 + sin_diff = bk.sin(diff) + # Mass matrix entries; only m12 is state-dependent. + m11 = (self.m1 + self.m2) * self.l1**2 + m12 = self.m2 * self.l1 * self.l2 * bk.cos(diff) + m22 = self.m2 * self.l2**2 + # b = tau - C omega - G, one (N, 1) column per joint. + b1 = ( + tau_1 + - self.m2 * self.l1 * self.l2 * sin_diff * omega_2 * omega_2 + - (self.m1 + self.m2) * self.g * self.l1 * bk.sin(theta_1) + ) + b2 = ( + tau_2 + + self.m2 * self.l1 * self.l2 * sin_diff * omega_1 * omega_1 + - self.m2 * self.g * self.l2 * bk.sin(theta_2) + ) + # theta_ddot = M^-1 b, 2x2 closed form. + det = m11 * m22 - m12 * m12 + theta_1_ddot = (b1 * m22 - m12 * b2) / det + theta_2_ddot = (m11 * b2 - b1 * m12) / det + + f = bk.stack([ + bk.ravel(omega_1), + bk.ravel(omega_2), + bk.ravel(theta_1_ddot), + bk.ravel(theta_2_ddot), + ]).T + return as_vector(bk, f, single) def get_model(self, x0=None, u0=None, eps=1e-6): """Get the discrete-time state-space model around an operating point. diff --git a/src/shinro/plants/inverted_pendulum.py b/src/shinro/plants/inverted_pendulum.py index 101bcf6..940d561 100644 --- a/src/shinro/plants/inverted_pendulum.py +++ b/src/shinro/plants/inverted_pendulum.py @@ -5,6 +5,7 @@ from shinro.components import PhysicsEngine, Plant from shinro.factories.registry import register_plant, register_plant_detector from shinro.utils.array_backend import ArrayBackend, NumpyBackend +from shinro.utils.batching import as_batch, as_vector, column, control_batch from shinro.utils.config_spec import BoundsConfig, strict_from_dict, strip_runtime_keys from shinro.utils.linearization import discretize_euler, linearize_plant @@ -148,20 +149,31 @@ def get_model(self, x0=None, u0=None, eps=1e-6): A_c, B_c = linearize_plant(self, x0, u0, eps=eps) return discretize_euler(A_c, B_c, self.dt, backend=self.bk) - def dynamics(self, state, control): + def dynamics(self, state, control, bk=None): """Continuous-time dynamics :math:`\\dot{x} = f(x, u)`. Args: - state: State vector (2,) — [theta, theta_dot]. - control: Control vector (1,) or scalar — [tau]. + state: State vector (2,) — [theta, theta_dot] — or a batch (N, 2). + control: Control (1,), batch (N, 1), or scalar — [tau]. + bk: Backend to evaluate with. Defaults to the plant's backend. Returns: - Time derivative of the state (2,) — [theta_dot, theta_ddot]. + Time derivative with the rank of ``state`` — + [theta_dot, theta_ddot]. """ - theta, theta_dot = state[0], state[1] - tau = control[0] if hasattr(control, '__len__') else control - theta_ddot = (self.g / self.l) * self.bk.sin(theta) + tau / (self.m * self.l**2) - (self.b / (self.m * self.l**2)) * theta_dot - return self.bk.stack([theta_dot, theta_ddot]) + bk = self.bk if bk is None else bk + x, single = as_batch(bk, state) + u = control_batch(bk, control, 1) + theta = column(bk, x, 0) + theta_dot = column(bk, x, 1) + tau = column(bk, u, 0) + theta_ddot = ( + (self.g / self.l) * bk.sin(theta) + + tau / (self.m * self.l**2) + - (self.b / (self.m * self.l**2)) * theta_dot + ) + f = bk.stack([bk.ravel(theta_dot), bk.ravel(theta_ddot)]).T + return as_vector(bk, f, single) def step(self, u): """Execute one control step. @@ -182,11 +194,9 @@ def step(self, u): self.state = self.get_state() return self.state - theta, theta_dot = self.state[0], self.state[1] - tau = u[0] if hasattr(u, '__len__') else u - theta_ddot = (self.g / self.l) * self.bk.sin(theta) + tau / (self.m * self.l**2) - (self.b / (self.m * self.l**2)) * theta_dot - theta_dot_new = theta_dot + theta_ddot * self.dt - theta_new = theta + theta_dot_new * self.dt + f = self.dynamics(self.state, u) + theta_dot_new = self.state[1] + f[1] * self.dt + theta_new = self.state[0] + theta_dot_new * self.dt self.state = self.bk.array([theta_new, theta_dot_new]) if self.state_bounds is not None: self.state = self.bk.clip(self.state, self.state_bounds[0], self.state_bounds[1]) diff --git a/src/shinro/runtime/graph_data.zig b/src/shinro/runtime/graph_data.zig index b8faea4..42681a8 100644 --- a/src/shinro/runtime/graph_data.zig +++ b/src/shinro/runtime/graph_data.zig @@ -6,6 +6,7 @@ pub const Op = enum { transpose, inv, reshape, clip, where_op, any, copy, tanh, relu, exp, argmax, one_hot, slice, sin, cos, stack, solve_qp, + abs, sign, pow, lt, min, }; pub const Node = struct { diff --git a/src/shinro/runtime/graph_data_manifest.json b/src/shinro/runtime/graph_data_manifest.json index 58e05d5..af37067 100644 --- a/src/shinro/runtime/graph_data_manifest.json +++ b/src/shinro/runtime/graph_data_manifest.json @@ -1,28 +1,37 @@ { + "buf_bytes": 2544, "buf_len": 318, + "clip_blob_bytes": 48, + "clip_blob_len": 3, + "const_blob_bytes": 1008, "const_blob_len": 126, "float_type": "f64", "has_solve_qp": false, + "input_bytes": 168, "inputs": [ { + "bytes": 24, "name": "y", "shape": [ 3 ] }, { + "bytes": 24, "name": "x_ref", "shape": [ 3 ] }, { + "bytes": 24, "name": "u_prev", "shape": [ 3 ] }, { + "bytes": 24, "name": "state_x_hat", "shape": [ 3, @@ -30,6 +39,7 @@ ] }, { + "bytes": 72, "name": "state_P", "shape": [ 3, @@ -646,25 +656,21 @@ "reshape", "sub" ], + "output_bytes": 24, "outputs": [ { + "bytes": 24, "name": "u", "shape": [ 3 ] } ], - "provenance": { - "configs": { - "configs/controllers/lqr_base.toml": "6257b61387fc3a2bd3a150f070ce02e15807921e2686eb0b7b99bb27b48e93a0", - "configs/estimators/kalman_base.toml": "693a1a66de459581014b981b35e5ed634ecf84223f67c8a2f89b05c94d61b428" - }, - "numpy_version": "2.4.6", - "python_version": "3.12.3" - }, "solve_qp": null, + "state_bytes": 120, "state_outputs": [ { + "bytes": 24, "name": "state_x_hat", "shape": [ 3, @@ -672,6 +678,7 @@ ] }, { + "bytes": 72, "name": "state_P", "shape": [ 3, @@ -679,6 +686,7 @@ ] }, { + "bytes": 24, "name": "state_u_prev", "shape": [ 3 diff --git a/src/shinro/runtime/linalg.zig b/src/shinro/runtime/linalg.zig index 345d75b..59de655 100644 --- a/src/shinro/runtime/linalg.zig +++ b/src/shinro/runtime/linalg.zig @@ -245,6 +245,49 @@ pub fn argmax (comptime m: usize, a:[]const f64) usize { return best; } +//minimum reductions--> numpy's np.min with axis=None / 0 / 1 + +/// Minimum over the whole flat array. numpy propagates NaN, so a NaN operand +/// wins the comparison and poisons the result — mirror that (a plain +/// `if (v < best)` would silently ignore NaNs, diverging from the oracle on +/// NaN feeds). +pub fn min_all (comptime m: usize, a: []const f64) [1]f64 { + var best: f64 = a[0]; + for (1..m) |i| { + const v = a[i]; + if (std.math.isNan(v) or v < best) best = v; + } + return .{best}; +} + +/// Minimum down each column: (rows, cols) -> (cols,), numpy `min(axis=0)`. +pub fn min_axis0 (comptime rows: usize, comptime cols: usize, a: []const f64) [cols]f64 { + var out: [cols]f64 = undefined; + for (0..cols) |j| { + var best = a[j]; + for (1..rows) |i| { + const v = a[i * cols + j]; + if (std.math.isNan(v) or v < best) best = v; + } + out[j] = best; + } + return out; +} + +/// Minimum across each row: (rows, cols) -> (rows,), numpy `min(axis=1)`. +pub fn min_axis1 (comptime rows: usize, comptime cols: usize, a: []const f64) [rows]f64 { + var out: [rows]f64 = undefined; + for (0..rows) |i| { + var best = a[i * cols]; + for (1..cols) |j| { + const v = a[i * cols + j]; + if (std.math.isNan(v) or v < best) best = v; + } + out[i] = best; + } + return out; +} + //one hot pub fn onehot (comptime depth: usize, idx:usize) [depth]f64{ var out: [depth]f64= undefined; diff --git a/src/shinro/runtime/lower.zig b/src/shinro/runtime/lower.zig index 157afad..3b123f5 100644 --- a/src/shinro/runtime/lower.zig +++ b/src/shinro/runtime/lower.zig @@ -96,10 +96,23 @@ export fn shinro_step(inputs: [*]const f64, outputs: [*]f64, state_out: [*]f64) .mul => ew2(g.nodes[0..], node, i, &buf, .mul), .div => ew2(g.nodes[0..], node, i, &buf, .div), .ne => ew2(g.nodes[0..], node, i, &buf, .ne), + .lt => ew2(g.nodes[0..], node, i, &buf, .lt), + .pow => ew2(g.nodes[0..], node, i, &buf, .pow), .neg => { const s = node_input(g.nodes[0..], node, &buf); inline for (0..node.rows * node.cols) |j| out[j] = -s[j]; }, + .abs => { + const s = node_input(g.nodes[0..], node, &buf); + inline for (0..node.rows * node.cols) |j| out[j] = @abs(s[j]); + }, + .sign => { + // Matches np.sign: -1 / 0 / +1 (0 maps to 0, not +1). + const s = node_input(g.nodes[0..], node, &buf); + inline for (0..node.rows * node.cols) |j| { + out[j] = if (s[j] > 0.0) 1.0 else if (s[j] < 0.0) -1.0 else 0.0; + } + }, .transpose => { const s = node_input(g.nodes[0..], node, &buf); // True 2-D transpose: out (node.rows, node.cols) = src.T, so @@ -192,6 +205,24 @@ export fn shinro_step(inputs: [*]const f64, outputs: [*]f64, state_out: [*]f64) const idx = la.argmax(n_in, s); out[0] = @floatFromInt(idx); }, + .min => { + // Minimum reduction; aux selects numpy's axis (0 = None / full, + // 1 = axis 0 down columns, 2 = axis 1 across rows). The output + // shape was fixed at trace time, so rows*cols is the exact + // element count the chosen reduction produces. + const s = node_input(g.nodes[0..], node, &buf); + const src = g.nodes[node.inputs[0]]; + if (node.aux == 0) { + const r = la.min_all(src.rows * src.cols, s); + out[0] = r[0]; + } else if (node.aux == 1) { + const r = la.min_axis0(src.rows, src.cols, s); + inline for (0..node.rows * node.cols) |j| out[j] = r[j]; + } else { + const r = la.min_axis1(src.rows, src.cols, s); + inline for (0..node.rows * node.cols) |j| out[j] = r[j]; + } + }, .one_hot => { const s = node_input(g.nodes[0..], node, &buf); const idx: usize = @intFromFloat(s[0]); @@ -253,7 +284,7 @@ export fn shinro_step(inputs: [*]const f64, outputs: [*]f64, state_out: [*]f64) // --- helpers --------------------------------------------------------------- -const BinOp = enum { add, sub, mul, div, ne }; +const BinOp = enum { add, sub, mul, div, ne, lt, pow }; /// Flat index of operand element (i, j) under numpy broadcasting. /// @@ -284,13 +315,13 @@ inline fn bcast_flat(op: g.Node, op_vec: bool, out_r: usize, out_c: usize, i: us /// node: The current add/sub/mul/div/ne node. /// self_idx: The node's index in `nodes` (its buffer offset). /// buf: The shared step buffer (written at the node's offset). -/// op: Which binary op to apply (add, sub, mul, div, ne). +/// op: Which binary op to apply (add, sub, mul, div, ne, lt, pow). inline fn ew2(nodes: []const g.Node, node: g.Node, self_idx: usize, buf: *[g.buf_len]f64, op: BinOp) void { const a = node_input_at(nodes, node.inputs[0], buf); const b = node_input_at(nodes, node.inputs[1], buf); const a_n = nodes[node.inputs[0]]; const b_n = nodes[node.inputs[1]]; - var out = buf.*[g.offsets[self_idx] ..][0 .. node.rows * node.cols]; + var out = buf.*[g.offsets[self_idx]..][0 .. node.rows * node.cols]; inline for (0..node.rows) |i| { inline for (0..node.cols) |j| { const av = a[bcast_flat(a_n, a_n.vec, node.rows, node.cols, i, j)]; @@ -303,6 +334,13 @@ inline fn ew2(nodes: []const g.Node, node: g.Node, self_idx: usize, buf: *[g.buf // Inequality as a 1.0/0.0 flag — the graph's boolean repr, // consumed by where_op downstream (e.g. PID anti-windup). .ne => if (av != bv) 1.0 else 0.0, + // Ordered comparison as a 1.0/0.0 flag — `ne`'s sibling, the + // predicate behind threshold guards (e.g. SMC's near-zero + // |c^T g| controllability check). + .lt => if (av < bv) 1.0 else 0.0, + // numpy's power semantics (np.power); SMC raises the abs'd + // sliding variable to a fractional alpha, so no negative base. + .pow => std.math.pow(f64, av, bv), }; } } diff --git a/src/shinro/utils/array_backend.py b/src/shinro/utils/array_backend.py index 62c435f..d319436 100644 --- a/src/shinro/utils/array_backend.py +++ b/src/shinro/utils/array_backend.py @@ -33,6 +33,22 @@ class ArrayBackend(ABC): identically in numpy and torch for 2D arrays. """ + def emit_named_output(self, name: str, value: Any) -> None: + """Publish an auxiliary named signal from a ``compute``/``estimate`` call. + + A concrete no-op for eager backends (numpy/torch): the signal is + already available to the caller, so nothing to publish. The tracing + backend overrides this to record a named graph ``output`` port, which + is how a component exposes a diagnostic alongside its primary return + value without changing its return contract (e.g. SMC's ``healthy`` + controllability flag). + + Args: + name: The output port name. + value: The signal (any array-like the backend produced). + """ + return None + @abstractmethod def array(self, data) -> Any: ... @@ -157,6 +173,9 @@ def slice_(self, x, start, stop) -> Any: ... @abstractmethod def sum(self, x, axis=None) -> Any: ... + @abstractmethod + def min(self, x, axis=None) -> Any: ... + @abstractmethod def reshape(self, x, *shape) -> Any: ... @@ -362,6 +381,9 @@ def slice_(self, x, start, stop): def sum(self, x, axis=None): return np.sum(x, axis=axis) + def min(self, x, axis=None): + return np.min(x, axis=axis) + def reshape(self, x, *shape): if len(shape) == 1 and isinstance(shape[0], (tuple, list)): shape = tuple(shape[0]) @@ -532,6 +554,13 @@ def trace(self, x): def clip(self, x, lo, hi): if not isinstance(x, self.torch.Tensor): x = self.torch.tensor(x, dtype=self.torch.float64) + # torch.clamp rejects numpy bounds (it takes a Number or a Tensor), so + # convert them. Array bounds are the common case (per-channel control + # limits), and they broadcast against the batched operand. + if lo is not None and not isinstance(lo, self.torch.Tensor): + lo = self.torch.as_tensor(lo, dtype=self.torch.float64, device=self.device) + if hi is not None and not isinstance(hi, self.torch.Tensor): + hi = self.torch.as_tensor(hi, dtype=self.torch.float64, device=self.device) return self.torch.clamp(x, lo, hi) def where(self, cond, a, b): @@ -601,6 +630,13 @@ def slice_(self, x, start, stop): def sum(self, x, axis=None): return self.torch.sum(x, dim=axis) + def min(self, x, axis=None): + # torch.min with a dim returns (values, indices); take the values so the + # signature matches numpy's min (which the interpreter/numpy backends use). + if axis is None: + return self.torch.min(x) + return self.torch.min(x, dim=axis).values + def reshape(self, x, *shape): if len(shape) == 1 and isinstance(shape[0], (tuple, list)): shape = tuple(shape[0]) diff --git a/src/shinro/utils/batched_adapter.py b/src/shinro/utils/batched_adapter.py index fd2bd78..58104b5 100644 --- a/src/shinro/utils/batched_adapter.py +++ b/src/shinro/utils/batched_adapter.py @@ -1,28 +1,40 @@ -"""Batched dynamics and cost adapter — adapts a Plant's single-state interface to batched ``(N, ...)`` arrays. +"""Batched dynamics and cost adapter — batched ``(N, ...)`` callables for sampling controllers. Sampling-based controllers (MPPI, CEM, iLQR, particle filters) roll out -:math:`N` parallel trajectories over a prediction horizon. The :class:`Plant` -interface is single-state: ``dynamics(x, u)`` returns :math:`\\dot{x}` for one -state, and ``get_model()`` returns ``(A, B)`` for one system. This adapter -bridges the two by exposing batched ``dynamics_fn(x_batch, u_batch, dt)`` and -``cost_fn(x_batch, u_batch, Q, R, x_ref=None)`` callables that operate on a -leading batch dimension :math:`N`. +:math:`N` parallel trajectories over a prediction horizon. This adapter +exposes the two callables they need — ``dynamics_fn(x_batch, u_batch, dt)`` +and ``cost_fn(x_batch, u_batch, Q, R, x_ref=None)`` — operating on a +leading batch dimension :math:`N`, from a plant's ``dynamics`` and +``get_model()``. Two dynamics paths are supported, dispatched on what the plant exposes: -* **LTI path** — plants whose ``get_model()`` returns ``(A, B)``. The state - update is a single batched matmul :math:`x_{k+1} = x_k A^T + u_k B^T`, - which runs as one native kernel on both numpy and torch. +* **Nonlinear path** — plants that implement ``dynamics`` (the ABC default is + ``None``, which linear plants keep). The plant evaluates its own + batch-capable derivative — one ``sin``/``mul`` node of shape ``(N, 1)`` per + formula term, not ``N`` private copies — and this adapter integrates it with + explicit Euler, ``x_{k+1} = x_k + dt f(x_k, u_k)``. Because the whole batch + rides in the node shapes, the traced graph's node count is independent of + the number of samples. The path is trace-safe: the plant routes every + backend call through the ``bk`` it is handed, because ``trace_node`` swaps + only the traced component's ``self.bk`` and leaves the plant's backend + concrete. + +* **LTI path** — plants whose ``dynamics`` returns ``None`` and whose + ``get_model()`` returns ``(A, B)``. The state update is a single batched + matmul ``x_{k+1} = x_k A^T + u_k B^T``, one native kernel on numpy and torch + alike. + +There is deliberately one nonlinear implementation, not two: a scalar formula +plus a batched rewrite would be two transcriptions of the same physics, free +to drift. ``Plant.dynamics`` accepts a single state ``(n_x,)`` **or** a batch +``(N, n_x)`` — a single state is a batch of one — so the eager per-sample +rollout, the finite-difference linearization, and the lowered graph all run +the same function. + +The adapter evaluates with the backend the caller passes (the component's +backend when tracing), defaulting to the plant's. -* **Nonlinear path** — plants that implement ``dynamics(x, u)``. The update is - a semi-implicit Euler step :math:`x_{k+1} = x_k + dt\\, f(x_k, u_k)`. - On torch the single-state dynamics is auto-vectorized over the batch with - ``torch.vmap``; on numpy a Python loop over the batch is used. - -The adapter uses the plant's own ``ArrayBackend`` throughout, so numpy and -torch both work with batched tensors and no hard-coded numpy in the hot loop. -For torch backends, the LTI path is a true batched matmul and the nonlinear -path runs through ``torch.vmap`` — both leverage torch's native batched ops. Args: plant: A :class:`Plant` instance exposing ``get_model()`` and (for the @@ -40,14 +52,13 @@ class BatchedDynamicsAdapter: The adapter detects the dynamics path once at construction: - * If ``plant.dynamics`` returns ``None`` (the ABC default — LTI plants - need not override it), the LTI matmul path is used. - * Otherwise the nonlinear path is used, integrating ``plant.dynamics`` - with semi-implicit Euler. On torch this is vectorized with - ``torch.vmap``; on numpy it loops over the batch. + * ``plant.dynamics`` returning non-``None`` (the ABC default is ``None``) + — the nonlinear path, integrating the plant's batch-capable derivative + with explicit Euler. Trace-safe. + * Otherwise the LTI matmul path. - All arrays live in the plant's backend, so torch inputs stay torch - throughout and run as batched native ops. + Arrays live in the plant's backend, or in the backend passed to + :meth:`dynamics_fn` (tracing passes the component's ``TraceBackend``). Usage: adapter = BatchedDynamicsAdapter(plant) @@ -66,12 +77,13 @@ def __init__(self, plant: Plant): self.D_x = self._A.shape[0] self.D_u = self._B.shape[1] - self._has_dynamics = plant.dynamics(state=self.bk.zeros(self.D_x), control=self.bk.zeros(self.D_u)) is not None - self._vmap = None - torch = getattr(self.bk, "torch", None) - if self._has_dynamics and torch is not None: - # Vectorize the single-state nonlinear dynamics over the batch. - self._vmap = torch.vmap(plant.dynamics, in_dims=(0, 0)) + # Path dispatch: a plant that overrides ``dynamics`` (the ABC default + # is None) rolls out nonlinearly; otherwise the linearized (A, B) + # matmul is exact and cheaper. The probe doubles as the capability + # check the ABC used to leave to ``dynamics() is not None``. + probe = plant.dynamics(state=self.bk.zeros(self.D_x), control=self.bk.zeros(self.D_u)) + self._nonlinear = probe is not None + self._path = "nonlinear" if self._nonlinear else "lti" @property def state_dim(self) -> int: @@ -83,37 +95,39 @@ def control_dim(self) -> int: """Control dimension :math:`D_u`.""" return self.D_u - def dynamics_fn(self, x_batch, u_batch, dt: float): - """Batched dynamics update. - - Args: - x_batch: Batch of states (N, D_x). - u_batch: Batch of controls (N, D_u). - dt: Time step (s). + @property + def dynamics_path(self) -> str: + """Which dynamics path this adapter dispatches to. - Returns: - Batch of next states (N, D_x). + ``"nonlinear"`` (the plant's own batch-capable derivative, integrated + with explicit Euler) or ``"lti"`` (one batched matmul). """ - if self._has_dynamics: - return self._integrate(x_batch, u_batch, dt) - return x_batch @ self._A.T + u_batch @ self._B.T + return self._path - def _integrate(self, x_batch, u_batch, dt: float): - """Semi-implicit Euler integration of the plant's nonlinear dynamics. + def dynamics_fn(self, x_batch, u_batch, dt: float, bk: Any | None = None): + """Batched dynamics update. + + Dispatches to the path chosen at construction. ``bk`` is the backend + the plant's nonlinear ``dynamics`` evaluates with, defaulting to the + plant's own: tracing passes the component's ``TraceBackend`` here, + because ``trace_node`` swaps only the traced component's ``self.bk`` — + a plant calling ``self.bk.sin`` on a tracer would otherwise reach a + concrete backend and fail. The LTI path ignores ``bk`` (operator + arithmetic already routes through the tracer). Args: x_batch: Batch of states (N, D_x). u_batch: Batch of controls (N, D_u). dt: Time step (s). + bk: Backend for the nonlinear path; defaults to the plant's backend. Returns: Batch of next states (N, D_x). """ - if self._vmap is not None: - return x_batch + dt * self._vmap(x_batch, u_batch) - return x_batch + dt * self.bk.stack( - [self.plant.dynamics(x_batch[i], u_batch[i]) for i in range(x_batch.shape[0])] - ) + if self._nonlinear: + f = self.plant.dynamics(x_batch, u_batch, bk=self.bk if bk is None else bk) + return x_batch + dt * f + return x_batch @ self._A.T + u_batch @ self._B.T def cost_fn(self, x_batch, u_batch, Q, R, x_ref: Any | None = None): """Batched quadratic stage cost. @@ -147,6 +161,13 @@ def _quad_form(self, z, W) -> Any: :math:`\\sum_i W_i z_i^2`) or a full ``(D, D)`` matrix (batched matmul). Returns a per-sample vector of shape ``(N,)``. + The row sum is written as a contraction rather than + ``bk.sum(..., axis=1)``: a sum over an axis *is* a matmul identity, and + the operator form traces — ``@`` lifts the concrete ones operand into a + const node — whereas ``bk.sum`` dispatches through whichever backend + the *plant* holds and never sees a traced operand. Same reduction, no + new op in the VM. + Args: z: Batch of vectors (N, D). W: Diagonal (D,) or full (D, D) weight matrix. @@ -157,5 +178,12 @@ def _quad_form(self, z, W) -> Any: if W is None: return self.bk.zeros(z.shape[0]) if W.ndim == 1: - return self.bk.sum(z * z * W, axis=1) - return self.bk.sum(z * (z @ W.T), axis=1) + # Diagonal weights: (z*z) contracted with W is exactly + # Σ_i W_i z_i² — W plays the contraction vector's role, so no ones + # vector is needed (and no rank-differing broadcast, which the + # tracer's elementwise ops reject even though the VM supports it). + return (z * z) @ W + # Full W: Σ_j z_j (z Wᵀ)_j, a row-wise dot product — contracted with a + # ones vector (there is no matmul identity for that one without it). + ones = self.bk.array([1.0] * z.shape[-1]) + return (z * (z @ W.T)) @ ones diff --git a/src/shinro/utils/batching.py b/src/shinro/utils/batching.py new file mode 100644 index 0000000..5693718 --- /dev/null +++ b/src/shinro/utils/batching.py @@ -0,0 +1,102 @@ +"""Rank helpers for batch-capable plant dynamics. + +``Plant.dynamics`` is called two ways: + +* **Per-sample** — a single ``(n_x,)`` state. The eager nonlinear rollout, the + finite-difference linearization (``linearize_plant``), and ``step()`` all use + this form. +* **Whole-batch** — an ``(N, n_x)`` state, evaluated by the MPPI rollout. This + is the form that is lowered, and its graph cost must not grow with ``N``. + +These helpers make the rank handling explicit and identical in every plant, so +the physics body is written once, in the column idiom, with no branching: + + x, single = as_batch(bk, state) # (1, n) or (N, n), plus the flag + theta = column(bk, x, 2) # (N, 1) + ... + return as_vector(bk, f, single) # (n,) or (N, n) + +The column idiom — ``bk.slice_(x.T, j, j + 1).T`` — transposes the ``(N, D)`` +batch to ``(D, N)``, slices one coordinate of every sample, and transposes back +to an ``(N, 1)`` column. Scalar indexing (``state[j]``) cannot be used on the +batched path because the tracing backend's :class:`~shinro.codegen.tracing.Tracer` +has no ``__getitem__``. + +Every helper takes the backend explicitly: plant code must call the backend it +was handed, not ``self.bk``, because ``trace_node`` swaps only the traced +component's backend — never the plant's. +""" + +from typing import Any + +from shinro.utils.array_backend import ArrayBackend + + +def as_batch(bk: ArrayBackend, x: Any) -> tuple[Any, bool]: + """Promote a 1-D vector to a ``(1, n)`` row. + + Args: + bk: Backend to reshape with. + x: State vector ``(n,)`` or batch ``(N, n)``. + + Returns: + ``(x_2d, was_1d)`` — the batch form and whether ``x`` was a single + vector (so the derivative can be converted back). + """ + if getattr(x, "shape", None) is None: + x = bk.array(x) + if len(x.shape) == 1: + return bk.reshape(x, (1, x.shape[0])), True + return x, False + + +def as_vector(bk: ArrayBackend, f: Any, was_1d: bool) -> Any: + """Undo :func:`as_batch` for a derivative. + + Args: + bk: Backend to reshape with. + f: Derivative batch ``(N, n)``. + was_1d: The flag returned by :func:`as_batch`. + + Returns: + ``(n,)`` when the state was a single vector, else ``f`` unchanged. + """ + return bk.ravel(f) if was_1d else f + + +def column(bk: ArrayBackend, x: Any, j: int) -> Any: + """Column ``j`` of a 2-D ``(N, D)`` batch as an ``(N, 1)`` column. + + Args: + bk: Backend to slice with. + x: 2-D batch ``(N, D)``. + j: Column index (compile-time constant). + + Returns: + Column ``j``, shape ``(N, 1)``. + """ + return bk.slice_(x.T, j, j + 1).T + + +def control_batch(bk: ArrayBackend, control: Any, n_u: int) -> Any: + """Normalize a control to a 2-D ``(N, n_u)`` (or ``(1, n_u)``) row. + + Accepts a scalar (``n_u == 1``, or the first input with the rest zero for + ``n_u > 1`` — the historical scalar-control convention), a ``(n_u,)`` + vector, or an ``(N, n_u)`` batch. + + Args: + bk: Backend to build with. + control: Scalar, ``(n_u,)``, or ``(N, n_u)``. + n_u: Control dimension. + + Returns: + A 2-D control row/batch. + """ + if getattr(control, "shape", None) is None: + control = bk.array([control] + [0.0] * (n_u - 1)) + if len(control.shape) == 0: + control = bk.reshape(bk.array([control] + [0.0] * (n_u - 1)), (n_u,)) + if len(control.shape) == 1: + return bk.reshape(control, (1, control.shape[0])) + return control diff --git a/tests/integration/scenarios/mppi_compile.toml b/tests/integration/scenarios/mppi_compile.toml new file mode 100644 index 0000000..b74b11a --- /dev/null +++ b/tests/integration/scenarios/mppi_compile.toml @@ -0,0 +1,32 @@ +# MPPI closed-loop compile scenario: KF + MPPI on the holonomic base. +# +# Unlike mppi_base_tracking.toml (a MuJoCo sim scenario with a sim-backed +# [plant] name), this one is compile-ready: +# * [plant] type/config lets the compile pipeline build the plant and derive +# the model, which is what wires MPPI's dynamics/cost (attach_plant). +# * MPPI's sampled perturbations arrive as a free C-ABI `epsilon` port of +# shape (N, K*D_u), sample-major — the kernel does no sampling. The host +# draws the noise each tick and packs it; see the MPPI module docstring. +[scenario] +name = "mppi_compile" +description = "KF + MPPI on the holonomic base, with host-supplied perturbations" +dt = 0.02 +input_limits = { min = [-0.5, -0.5, -1.0], max = [0.5, 0.5, 1.0] } + +[plant] +type = "HolonomicMobileRobot" +config = "configs/plants/holonomic_base.toml" + +[controller] +type = "MPPI" +config = "configs/controllers/mppi_base.toml" + +[estimator] +type = "KalmanFilter" +config = "configs/estimators/kalman_base.toml" + +# Build spec: dims are baked into the graph at trace time. MPPI's N/K/D_u come +# from the controller config (and the plant's control dim), not from here. +[compile] +n_x = 3 +n_u = 3 diff --git a/tests/integration/scenarios/mppi_pendulum_compile.toml b/tests/integration/scenarios/mppi_pendulum_compile.toml new file mode 100644 index 0000000..efc8aff --- /dev/null +++ b/tests/integration/scenarios/mppi_pendulum_compile.toml @@ -0,0 +1,35 @@ +# MPPI closed-loop compile scenario: KF + MPPI on the inverted pendulum. +# +# The nonlinear counterpart of mppi_compile.toml. The plant is a real +# InvertedPendulum (not a sim-backed name), so the compile pipeline builds it, +# derives the KF's linearized (A, B) via get_model(), and wires MPPI's rollout +# through attach_plant — which is the plant's own batch-capable `dynamics`. +# That makes this the end-to-end proof that a nonlinear rollout lowers: the +# graph contains sin/mul nodes over the sample batch, and the same three-way +# oracle (numpy / interpreter / .so) checks it as for the LTI case. +# +# The perturbations still arrive as a free C-ABI `epsilon` port of shape +# (N, K*D_u): sampling stays on the host. +[scenario] +name = "mppi_pendulum_compile" +description = "KF + MPPI on the inverted pendulum, with a nonlinear rollout" +dt = 0.01 +input_limits = { min = [-0.5], max = [0.5] } + +[plant] +type = "InvertedPendulum" +config = "configs/plants/inverted_pendulum.toml" + +[controller] +type = "MPPI" +config = "configs/controllers/mppi_inverted_pendulum.toml" + +[estimator] +type = "KalmanFilter" +config = "configs/estimators/kalman_inverted_pendulum.toml" + +# Build spec: dims are baked into the graph at trace time. MPPI's N/K/D_u come +# from the controller config (and the plant's control dim), not from here. +[compile] +n_x = 2 +n_u = 1 diff --git a/tests/test_batched_adapter.py b/tests/test_batched_adapter.py index 2ca524d..cb55ddc 100644 --- a/tests/test_batched_adapter.py +++ b/tests/test_batched_adapter.py @@ -42,8 +42,13 @@ def test_lti_zero_input_keeps_state(self, bk): out = adapter.dynamics_fn(x, u, 0.02) assert np.allclose(_to_np(out, bk), _to_np(x, bk)) - def test_nonlinear_dynamics_matches_euler(self, bk): - """The nonlinear dynamics equals a manual Euler step of plant.dynamics.""" + def test_nonlinear_dynamics_matches_per_sample_calls(self, bk): + """The rollout equals a manual Euler step of per-sample ``dynamics`` calls. + + The plant's ``dynamics`` is batch-capable, so this checks the rank + handling: the ``(N, ·)`` call the adapter makes must agree with ``N`` + single-state calls. + """ from shinro.utils.batched_adapter import BatchedDynamicsAdapter plant = self._pendulum_plant(bk) adapter = BatchedDynamicsAdapter(plant) @@ -56,7 +61,7 @@ def test_nonlinear_dynamics_matches_euler(self, bk): x_next = x[i] + dt * plant.dynamics(x[i], u[i]) expected.append(_to_np(x_next, bk)) assert _to_np(out, bk).shape == (5, 2) - assert np.allclose(_to_np(out, bk), np.array(expected), atol=1e-10) + assert np.allclose(_to_np(out, bk), np.array(expected), atol=1e-12) def test_cost_matches_analytic(self, bk): """The batched cost equals (x-x_ref)ᵀQ(x-x_ref) + uᵀRu per sample.""" @@ -118,35 +123,79 @@ def test_torch_backend_keeps_tensors(self, bk): assert cost.shape == (4,) -class TestBatchedAdapterNonlinearVmap: - """Verify torch.vmap vectorizes nonlinear plant dynamics over the batch.""" +class TestDynamicsDispatch: + """Which dynamics path the adapter picks, and that the nonlinear one traces.""" - def test_vmap_pendulum(self, bk): - """vmap(pendulum.dynamics) produces (N, 2) from (N, 2), (N, 1).""" - torch = pytest.importorskip("torch") - if not hasattr(bk, "torch"): - pytest.skip("requires TorchBackend") + def _pendulum(self, bk): from shinro.plants.inverted_pendulum import InvertedPendulum - plant = InvertedPendulum(backend=bk) - x = bk.array([[0.1, 0.0], [0.5, 0.3], [1.0, -1.0], [0.0, 0.0], [-0.2, 0.7]]) - u = bk.array([[0.1], [0.2], [0.0], [0.5], [-0.3]]) - out = torch.vmap(plant.dynamics, in_dims=(0, 0))(x, u) - assert out.shape == (5, 2) - for i in range(5): - single = plant.dynamics(x[i], u[i]) - assert torch.allclose(out[i], single) + return InvertedPendulum(mass=0.1, length=0.5, damping=0.0, gravity=9.81, dt=0.01, backend=bk) - def test_vmap_cartpole(self, bk): - """vmap(cartpole.dynamics) produces (N, 4) from (N, 4), (N, 1).""" - torch = pytest.importorskip("torch") - if not hasattr(bk, "torch"): - pytest.skip("requires TorchBackend") - from shinro.plants.cartpole import CartPole - plant = CartPole(backend=bk) - x = bk.array([[0.0, 0.0, 0.1, 0.0], [0.0, 0.0, 0.5, 0.0], [0.0, 0.0, 1.0, 0.0]]) - u = bk.array([[1.0], [0.0], [-1.0]]) - out = torch.vmap(plant.dynamics, in_dims=(0, 0))(x, u) - assert out.shape == (3, 4) - for i in range(3): - single = plant.dynamics(x[i], u[i]) - assert torch.allclose(out[i], single) + def test_nonlinear_vs_lti(self, bk): + """A plant with ``dynamics`` is nonlinear; an (A, B)-only plant is LTI.""" + from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + from shinro.utils.batched_adapter import BatchedDynamicsAdapter + assert BatchedDynamicsAdapter(self._pendulum(bk)).dynamics_path == "nonlinear" + base = HolonomicMobileRobot( + num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=0.02, backend=bk + ) + assert BatchedDynamicsAdapter(base).dynamics_path == "lti" + + def test_nonlinear_path_uses_caller_backend(self): + """Tracers plus a caller-supplied backend: the plant emits graph nodes.""" + from shinro.codegen.trace_backend import TraceBackend + from shinro.codegen.tracing import Graph, Tracer + from shinro.utils.array_backend import NumpyBackend + from shinro.utils.batched_adapter import BatchedDynamicsAdapter + adapter = BatchedDynamicsAdapter(self._pendulum(NumpyBackend())) + g = Graph() + x = Tracer(g, (3, 2), g.input("x", (3, 2))) + u = Tracer(g, (3, 1), g.input("u", (3, 1))) + out = adapter.dynamics_fn(x, u, 0.01, bk=TraceBackend(g)) + assert isinstance(out, Tracer) + assert out.shape == (3, 2) + ops = [n.op for n in g.nodes] + assert "sin" in ops + assert "slice" in ops + assert "stack" in ops + + def test_attach_plant_routes_controller_backend(self, bk): + """attach_plant evaluates the nonlinear dynamics with the controller's bk.""" + from shinro.codegen.trace_backend import TraceBackend + from shinro.codegen.tracing import Graph, Tracer + from shinro.controllers.mppi import MPPIController + ctrl = MPPIController( + num_samples=3, temperature=1.0, dt=0.01, horizon=2, + noise_sigma=[0.5], u_min=[-1.0], u_max=[1.0], seed=0, backend=bk, + ) + ctrl.attach_plant(self._pendulum(bk)) + g = Graph() + x = Tracer(g, (3, 2), g.input("x", (3, 2))) + u = Tracer(g, (3, 1), g.input("u", (3, 1))) + original = ctrl.bk + ctrl.bk = TraceBackend(g) # type: ignore[assignment] # tracing swaps in a TraceBackend + try: + assert ctrl.dynamics_fn is not None + out = ctrl.dynamics_fn(x, u, 0.01) + finally: + ctrl.bk = original + assert isinstance(out, Tracer) + assert out.shape == (3, 2) + assert any(n.op == "sin" for n in g.nodes) + + def test_pendulum_dynamics_traces(self): + """InvertedPendulum.dynamics emits graph nodes under a TraceBackend.""" + from shinro.codegen.trace_backend import TraceBackend + from shinro.codegen.tracing import Graph, Tracer + from shinro.plants.inverted_pendulum import InvertedPendulum + from shinro.utils.array_backend import NumpyBackend + plant = InvertedPendulum(backend=NumpyBackend()) + g = Graph() + x = Tracer(g, (4, 2), g.input("x", (4, 2))) + u = Tracer(g, (4, 1), g.input("u", (4, 1))) + out = plant.dynamics(x, u, bk=TraceBackend(g)) + assert isinstance(out, Tracer) + assert out.shape == (4, 2) + ops = [n.op for n in g.nodes] + assert "sin" in ops + assert "slice" in ops + assert "stack" in ops diff --git a/tests/test_codegen_build.py b/tests/test_codegen_build.py index 97b9b8d..a5074dc 100644 --- a/tests/test_codegen_build.py +++ b/tests/test_codegen_build.py @@ -12,12 +12,15 @@ from __future__ import annotations import tomllib +from dataclasses import dataclass import numpy as np import pytest -from shinro.codegen import build_composed_graph -from shinro.codegen.compose import ComposedGraph +from shinro.codegen import build_composed_graph, interpret +from shinro.codegen.compose import ComposedGraph, _lookup_input_shape +from shinro.components import StateEstimator +from shinro.factories.registry import register_estimator # Input limits from base_tracking.toml's [scenario.input_limits]. _BASE_LIMITS = (np.array([-0.5, -0.5, -1.0]), np.array([0.5, 0.5, 1.0])) @@ -293,3 +296,166 @@ def test_derived_plant_scenario_matches_explicit(tmp_path): ) _assert_graphs_identical(cg_derived, cg_explicit) + + +def test_mppi_composes_with_a_host_filled_epsilon_port(): + """MPPI's sampled perturbations are a free host input, not an estimator feed. + + The compile path attaches the plant (MPPI's dynamics/cost come from its + model) and then composes; ``epsilon`` is declared with MPPI's own shape — + ``(N, K*D_u)``, sample-major — and appended last, so the + estimator/controller ports ahead of it are unchanged. + """ + from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + from shinro.utils.array_backend import NumpyBackend + + plant = HolonomicMobileRobot( + num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=0.02, backend=NumpyBackend() + ) + cg = build_composed_graph( + "configs/estimators/kalman_base.toml", + "configs/controllers/mppi_base.toml", + 3, + 3, + input_limits=_BASE_LIMITS, + plant=plant, + ) + + # The shipped port layout is unchanged; epsilon arrives last. + assert cg.inputs[:5] == ["y", "x_ref", "u_prev", "state_x_hat", "state_P"] + assert cg.inputs[-1] == "epsilon" + assert _lookup_input_shape(cg.graph, "epsilon", default=()) == (200, 15 * 3) + # MPPI's diagnostic survives composition, appended after the control output. + assert cg.outputs == ["u", "costs"] + + feeds = {name: np.zeros(_lookup_input_shape(cg.graph, name, default=())) for name in cg.inputs} + out = interpret(cg.graph, feeds) + assert np.asarray(out["u"]).shape == (3,) + assert np.asarray(out["costs"]).shape == (200,) + + +# ─── a hypothetical third estimator: does the pipeline generalize? ──────── + + +@dataclass(frozen=True) +class _LeakyAverageConfig: + """Strict config for the hypothetical estimator below.""" + + n_x: int + alpha: float = 0.5 + name: str = "leaky_average" + + +@register_estimator("LeakyAverage") +class _LeakyAverageEstimator(StateEstimator): + """A plausible *new* estimator: an exponential moving average of the measurement. + + Deliberately not shipped — it exists to prove the compile pipeline + generalizes past the two registered estimators. It has one recurrent attr + (``x_hat``) that no tracer code declares: ``build_composed_graph``'s + two-pass state discovery finds it by attr-diff and promotes it to a + ``state_x_hat`` port. ``estimate`` is operator/``bk`` only, so it traces. + """ + + Config = _LeakyAverageConfig + + def __init__(self, n_x: int, alpha: float = 0.5, backend=None): + from shinro.utils.array_backend import NumpyBackend + + self.bk = backend or NumpyBackend() + self.n_x = n_x + self.alpha = float(alpha) + self.x_hat = self.bk.zeros(n_x) + + def estimate(self, measurement, control_input): + y = self.bk.ravel(measurement) + self.x_hat = (1.0 - self.alpha) * self.x_hat + self.alpha * y + return self.bk.reshape(self.x_hat, (self.n_x, 1)) + + def reset(self): + self.x_hat = self.bk.zeros(self.n_x) + + @classmethod + def from_config(cls, config, backend=None): + cfg = cls.parse_config(config) + return cls(n_x=cfg.n_x, alpha=cfg.alpha, backend=backend) + + +def test_composes_with_a_hypothetical_third_estimator(): + """A brand-new estimator composes with MPPI and matches the live loop. + + Nothing in the pipeline is estimator-aware: the class is registered ad hoc, + its inputs follow the ``measurement`` / ``control_input`` naming contract, + and its undiscovered state is promoted automatically. The composed graph + must then run tick-for-tick like the live components, with MPPI's + host-supplied ``epsilon`` alongside the new estimator's state port. + """ + from shinro.codegen import interpret + from shinro.factories.controller_factory import ControllerFactory + from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + from shinro.utils.array_backend import NumpyBackend + + bk = NumpyBackend() + plant = HolonomicMobileRobot( + num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=0.02, backend=bk + ) + est_cfg = {"type": "LeakyAverage", "n_x": 3, "alpha": 0.4} + cg = build_composed_graph( + est_cfg, + "configs/controllers/mppi_base.toml", + 3, + 3, + input_limits=_BASE_LIMITS, + plant=plant, + ) + + # The new estimator's state became a recurrent port; epsilon is still free + # and still last, and the controller never learned about either change. + assert cg.inputs == ["y", "x_ref", "u_prev", "state_x_hat", "state_u", "epsilon"] + assert "state_x_hat" in cg.state_outputs + assert "state_u" in cg.state_outputs + + # Tick-for-tick against the live components (same order as the graph). + est = _LeakyAverageEstimator.from_config(est_cfg, backend=NumpyBackend()) + ctrl = ControllerFactory("configs/controllers/mppi_base.toml").create(backend=NumpyBackend()) + ctrl.attach_plant(plant) + + rng = np.random.default_rng(5) + n_x, n_u, n_samples, horizon = 3, 3, ctrl.N, ctrl.K + x_hat = np.zeros(n_x) + plan = np.zeros((horizon, n_u)) + u_prev = np.zeros(n_u) + + for trial in range(3): + y = rng.normal(0.0, 0.2, n_x) + x_ref = rng.normal(0.0, 0.5, n_x) + eps = rng.normal(0.0, 0.3, (n_samples, horizon * n_u)) + + traced = interpret( + cg.graph, + { + "y": y, + "x_ref": x_ref, + "u_prev": u_prev, + "state_x_hat": x_hat, + "state_u": plan, + "epsilon": eps, + }, + ) + + # Live reference: estimator, then controller, then the composed clip. + x_hat_next = np.asarray(est.estimate(y.reshape(-1, 1), u_prev.reshape(-1, 1))).ravel() + u_live = np.clip( + np.asarray(ctrl.compute(x_hat_next, x_ref, eps)), _BASE_LIMITS[0], _BASE_LIMITS[1] + ) + + np.testing.assert_allclose(traced["u"], u_live, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose( + np.asarray(traced["state_x_hat"]).ravel(), x_hat_next, rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose(np.asarray(traced["state_u"]), ctrl.u, rtol=1e-12, atol=1e-12) + + # Carry the live state forward; the graph's state outputs match it. + x_hat = x_hat_next + plan = ctrl.u.copy() + u_prev = u_live diff --git a/tests/test_codegen_compose.py b/tests/test_codegen_compose.py index 82d1115..a762078 100644 --- a/tests/test_codegen_compose.py +++ b/tests/test_codegen_compose.py @@ -614,6 +614,154 @@ def test_unknown_controller_input_raises(self): compose(estimator, controller, plant_dims=_BASE_DIMS, input_limits=None) +class TestComposeHostInputs: + """Host-filled free inputs (e.g. MPPI's sampled perturbations). + + A controller may declare inputs the *host* fills each tick rather than the + estimator or reference. They become composed-graph input ports whose shape + comes from the controller's own declaration, and — unlike state ports — + they are neither wired to the estimator nor fed back next tick. + """ + + @staticmethod + def _graphs(host_shape=(4, 6)): + from shinro.codegen.trace_node import NodeGraph + + est_g = Graph() + est_in = est_g.input("measurement", (3, 1)) + est_g.input("control_input", (3, 1)) + est_g.input("state_x_hat", (3, 1)) + est_g.output("out", est_in) + estimator = NodeGraph( + graph=est_g, + contract=None, # type: ignore[arg-type] + input_nodes={"measurement": est_in}, + output_nodes={"out": est_in}, + state_attrs=[], + ) + ctrl_g = Graph() + ctrl_in = ctrl_g.input("current_state", (3,)) + ctrl_g.input("target_state", (3,)) + ctrl_g.input("epsilon", host_shape) # free host input (not consumed here) + ctrl_g.output("out", ctrl_in) + controller = NodeGraph( + graph=ctrl_g, + contract=None, # type: ignore[arg-type] + input_nodes={"current_state": ctrl_in}, + output_nodes={"out": ctrl_in}, + state_attrs=[], + ) + return estimator, controller + + def test_host_input_becomes_a_free_port_appended_last(self): + estimator, controller = self._graphs() + cg = compose(estimator, controller, plant_dims=_BASE_DIMS, host_inputs=("epsilon",)) + + # Appended last, so every existing port layout is unchanged. + assert cg.inputs[-1] == "epsilon" + # Shape comes from the controller's own declaration (not (n_x,)). + assert _lookup_input_shape(cg.graph, "epsilon", default=()) == (4, 6) + # Free, not recurrent: never fed back, never wired to the estimator. + assert "epsilon" not in cg.state_inputs + assert "epsilon" not in cg.state_outputs + + def test_undeclared_host_input_still_raises(self): + """The role guard stays loud — a name must be declared to be free.""" + estimator, controller = self._graphs() + with pytest.raises(ValueError, match="does not map to a known role"): + compose(estimator, controller, plant_dims=_BASE_DIMS) + + def test_host_port_is_a_real_graph_input(self): + """The interpreter needs the free port fed, like any other input.""" + estimator, controller = self._graphs() + cg = compose(estimator, controller, plant_dims=_BASE_DIMS, host_inputs=("epsilon",)) + feeds = {name: np.zeros(_lookup_input_shape(cg.graph, name, default=())) for name in cg.inputs} + out = interpret(cg.graph, feeds) + assert np.asarray(out["u"]).shape == (3,) + + +class TestComposeDiagnostics: + """Auxiliary outputs (``emit_named_output`` diagnostics) survive composition. + + A composed binary used to expose only the control vector: the merge skipped + subgraph output markers, so MPPI's ``costs`` / SMC's ``healthy`` were + dropped. They are now forwarded, appended after ``u`` so existing port + layouts stay unchanged, and a duplicate name is a loud error (a silent + shadow would hide one component's signal). + """ + + @staticmethod + def _graphs(est_diag: str | None = "estimator_ok", ctrl_diag: str | None = "healthy"): + from shinro.codegen.trace_node import NodeGraph + + est_g = Graph() + est_in = est_g.input("measurement", (3, 1)) + est_g.input("control_input", (3, 1)) + est_g.input("state_x_hat", (3, 1)) + est_g.output("out", est_in) + est_outs = {"out": est_in} + if est_diag: + est_g.output(est_diag, est_in) + est_outs[est_diag] = est_in + estimator = NodeGraph( + graph=est_g, + contract=None, # type: ignore[arg-type] + input_nodes={"measurement": est_in}, + output_nodes=est_outs, + state_attrs=[], + ) + + ctrl_g = Graph() + ctrl_in = ctrl_g.input("current_state", (3,)) + ctrl_g.input("target_state", (3,)) + flag = ctrl_g.emit("const", [], (), value=np.float64(1.0)) + ctrl_g.output("out", ctrl_in) + ctrl_outs = {"out": ctrl_in} + if ctrl_diag: + ctrl_g.output(ctrl_diag, flag) + ctrl_outs[ctrl_diag] = flag + controller = NodeGraph( + graph=ctrl_g, + contract=None, # type: ignore[arg-type] + input_nodes={"current_state": ctrl_in}, + output_nodes=ctrl_outs, + state_attrs=[], + ) + return estimator, controller + + def test_diagnostics_forwarded_after_u(self): + estimator, controller = self._graphs() + cg = compose(estimator, controller, plant_dims=_BASE_DIMS) + + # Appended after the control output; estimator's before the controller's. + assert cg.outputs == ["u", "estimator_ok", "healthy"] + # Not recurrent, and not confused with the control vector. + assert "healthy" not in cg.state_outputs + + def test_diagnostic_is_a_real_output_port(self): + estimator, controller = self._graphs() + cg = compose(estimator, controller, plant_dims=_BASE_DIMS) + feeds = {name: np.zeros(_lookup_input_shape(cg.graph, name, default=())) for name in cg.inputs} + out = interpret(cg.graph, feeds) + assert float(np.asarray(out["healthy"]).reshape(-1)[0]) == 1.0 + assert "estimator_ok" in out + + def test_no_diagnostics_leaves_the_layout_unchanged(self): + estimator, controller = self._graphs(est_diag=None, ctrl_diag=None) + cg = compose(estimator, controller, plant_dims=_BASE_DIMS) + assert cg.outputs == ["u"] + + def test_duplicate_diagnostic_names_raise(self): + estimator, controller = self._graphs(est_diag="status", ctrl_diag="status") + with pytest.raises(ValueError, match="published by both"): + compose(estimator, controller, plant_dims=_BASE_DIMS) + + def test_diagnostic_colliding_with_the_control_output_raises(self): + estimator, controller = self._graphs(ctrl_diag="u") + with pytest.raises(ValueError, match="published by both"): + compose(estimator, controller, plant_dims=_BASE_DIMS) + + # ─── Test 10: compose error paths ────────────────────────────────────────── diff --git a/tests/test_controllers.py b/tests/test_controllers.py index f02bf47..b6f68ec 100644 --- a/tests/test_controllers.py +++ b/tests/test_controllers.py @@ -4,7 +4,7 @@ def _to_np(x, bk): """Convert a backend array to numpy for assertion comparisons.""" - return bk.to_numpy(x) if hasattr(bk, 'to_numpy') else x + return bk.to_numpy(x) if hasattr(bk, "to_numpy") else x class TestLQR: @@ -17,6 +17,7 @@ def test_lqr_gain_stabilizes_1d(self, bk): Q = bk.eye(1) R = bk.eye(1) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) K = lqr.K A_cl = A - B @ K @@ -30,6 +31,7 @@ def test_lqr_gain_analytical_1d(self, bk): Q = bk.eye(1) R = bk.eye(1) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) K = _to_np(lqr.K, bk)[0, 0] P_expected = (1 + np.sqrt(5)) / 2 @@ -43,8 +45,10 @@ def test_lqr_dare_residual(self, bk): Q = bk.eye(1) R = bk.eye(1) from shinro.controllers.lqr import LQR + LQR(Q, R, A, B, backend=bk) from scipy.linalg import solve_discrete_are + P = solve_discrete_are(_to_np(A, bk), _to_np(B, bk), _to_np(Q, bk), _to_np(R, bk)) residual = A.T @ P @ A - P - A.T @ P @ B @ np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A) + Q assert np.linalg.norm(residual) < 1e-10 @@ -56,8 +60,10 @@ def test_lqr_gain_formula(self, bk): Q = bk.eye(2) R = 0.5 * bk.eye(2) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) from scipy.linalg import solve_discrete_are + P_np = solve_discrete_are(_to_np(A, bk), _to_np(B, bk), _to_np(Q, bk), _to_np(R, bk)) P = bk.from_numpy(P_np) K_expected = bk.inv(R + B.T @ P @ B) @ (B.T @ P @ A) @@ -70,6 +76,7 @@ def test_lqr_closed_loop_eigenvalues(self, bk): Q = bk.eye(2) R = bk.array([[0.1]]) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) A_cl = A - B @ lqr.K eigs = np.linalg.eigvals(_to_np(A_cl, bk)) @@ -82,6 +89,7 @@ def test_lqr_optimal_control_law(self, bk): Q = bk.eye(2) R = bk.eye(2) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) x = bk.array([1.5, -0.7]) u = lqr.compute(x) @@ -95,6 +103,7 @@ def test_lqr_compute_shape(self, bk): Q = bk.eye(2) R = bk.eye(2) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) x = bk.array([1.0, 2.0]) u = lqr.compute(x) @@ -107,6 +116,7 @@ def test_lqr_regulation_to_zero(self, bk): Q = bk.eye(2) R = bk.eye(2) from shinro.controllers.lqr import LQR + lqr = LQR(Q, R, A, B, backend=bk) x = bk.array([1.0, 2.0]) u = lqr.compute(x) @@ -117,6 +127,7 @@ def test_lqr_from_config(self, bk): """from_config creates a valid LQR controller with a gain matrix.""" config = {"state_cost": [1.0, 1.0], "control_cost": [1.0, 1.0], "dt": 0.1} from shinro.controllers.lqr import LQR + lqr = LQR.from_config(config, backend=bk) assert lqr.K is not None @@ -130,6 +141,7 @@ def test_lqr_from_config_full_matrix(self, bk): "dt": 0.1, } from shinro.controllers.lqr import LQR + lqr = LQR.from_config(config, backend=bk) assert lqr.K is not None assert _to_np(lqr.Q, bk).shape == (2, 2) @@ -143,6 +155,7 @@ class TestPID: def test_pid_derivative_zero_on_first_call(self, bk): """Derivative term is zero on the first call (no previous error).""" from shinro.controllers.pid import PIDController + pid = PIDController( kp=bk.array([1.0]), ki=bk.array([0.0]), @@ -157,6 +170,7 @@ def test_pid_derivative_zero_on_first_call(self, bk): def test_pid_derivative_on_second_call(self, bk): """Derivative term on the second call is kd * (e_k - e_{k-1}) / dt.""" from shinro.controllers.pid import PIDController + pid = PIDController( kp=bk.array([0.0]), ki=bk.array([0.0]), @@ -172,6 +186,7 @@ def test_pid_derivative_on_second_call(self, bk): def test_pid_integral_accumulates(self, bk): """The integral term accumulates error over successive calls.""" from shinro.controllers.pid import PIDController + pid = PIDController( kp=bk.array([0.0]), ki=bk.array([1.0]), @@ -189,6 +204,7 @@ def test_pid_integral_accumulates(self, bk): def test_pid_output_limits_clamp(self, bk): """Output limits clamp the control effort to [min, max].""" from shinro.controllers.pid import PIDController + pid = PIDController( kp=bk.array([10.0]), ki=bk.array([0.0]), @@ -203,6 +219,7 @@ def test_pid_output_limits_clamp(self, bk): def test_pi_eliminates_steady_state_error(self, bk): """PI control drives a first-order lag plant to the target with zero steady-state error.""" from shinro.controllers.pid import PIDController + # First-order lag: x_{k+1} = a*x + b*dt*u, tau=0.1s, dt=0.01s. a = float(np.exp(-0.01 / 0.1)) b = 1.0 @@ -223,6 +240,7 @@ def test_pi_eliminates_steady_state_error(self, bk): def test_p_only_steady_state_error(self, bk): """P-only control leaves a non-zero steady-state error for a first-order lag plant.""" from shinro.controllers.pid import PIDController + # First-order lag: x_{k+1} = a*x + b*dt*u, tau=0.1s, dt=0.01s. a = float(np.exp(-0.01 / 0.1)) b = 1.0 @@ -248,6 +266,7 @@ def test_p_only_steady_state_error(self, bk): def test_pid_anti_windup(self, bk): """When output is clamped, the integral term back-calculates on saturated channels.""" from shinro.controllers.pid import PIDController + lo = bk.array([-0.5]) hi = bk.array([0.5]) pid = PIDController( @@ -268,6 +287,7 @@ def test_pid_anti_windup(self, bk): def test_pid_reset(self, bk): """reset() clears the integral accumulator and previous error.""" from shinro.controllers.pid import PIDController + pid = PIDController( kp=bk.array([1.0]), ki=bk.array([1.0]), @@ -285,6 +305,7 @@ def test_pid_from_config(self, bk): """from_config creates a valid PID controller.""" config = {"kp": [1.0], "ki": [0.5], "kd": [0.1], "dt": 0.01} from shinro.controllers.pid import PIDController + pid = PIDController.from_config(config, backend=bk) assert pid.kp is not None @@ -295,6 +316,7 @@ class TestMPC: def test_mpc_H_symmetric(self, bk): """The QP Hessian H is symmetric.""" from shinro.controllers.mpc_lti import MPC_LTI + n = 2 m = 2 A = bk.eye(n) @@ -302,14 +324,14 @@ def test_mpc_H_symmetric(self, bk): Q = bk.eye(n) R = bk.eye(m) P = bk.eye(n) - mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, - A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) + mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) H = _to_np(mpc.H, bk) assert np.allclose(H, H.T) def test_mpc_F_shape(self, bk): """The QP linear term F has shape (n_x, N * n_u).""" from shinro.controllers.mpc_lti import MPC_LTI + n = 2 m = 2 A = bk.eye(n) @@ -317,14 +339,14 @@ def test_mpc_F_shape(self, bk): Q = bk.eye(n) R = bk.eye(m) P = bk.eye(n) - mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, - A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) + mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) F = _to_np(mpc.F, bk) assert F.shape == (n, 5 * m) def test_mpc_compute_shape(self, bk): """compute() returns a control vector of dimension n_u.""" from shinro.controllers.mpc_lti import MPC_LTI + n = 2 m = 2 A = bk.eye(n) @@ -332,8 +354,7 @@ def test_mpc_compute_shape(self, bk): Q = bk.eye(n) R = bk.eye(m) P = bk.eye(n) - mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, - A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) + mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) F = bk.eye(m) mpc.constraints(F, bk.array([1.0, 1.0]), bk.array([-1.0, -1.0])) x0 = bk.array([1.0, 0.0]) @@ -343,6 +364,7 @@ def test_mpc_compute_shape(self, bk): def test_mpc_constraints_respected(self, bk): """MPC respects hard input constraints |u| <= bound.""" from shinro.controllers.mpc_lti import MPC_LTI + n = 2 m = 2 A = bk.eye(n) @@ -350,8 +372,7 @@ def test_mpc_constraints_respected(self, bk): Q = bk.eye(n) R = bk.eye(m) P = bk.eye(n) - mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, - A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) + mpc = MPC_LTI(horizon=5, control_cost_matrix=R, state_cost_matrix=Q, A_dynamics=A, B_dynamics=B, terminal_cost=P, backend=bk) bound = 0.5 F = bk.eye(m) mpc.constraints(F, bk.array([bound, bound]), bk.array([-bound, -bound])) @@ -369,6 +390,7 @@ def test_mpc_from_config(self, bk): "dt": 0.1, } from shinro.controllers.mpc_lti import MPC_LTI_Base + mpc = MPC_LTI_Base.from_config(config, backend=bk) assert mpc.H is not None assert mpc.F is not None @@ -384,6 +406,7 @@ def test_mpc_from_config_full_matrix(self, bk): "dt": 0.1, } from shinro.controllers.mpc_lti import MPC_LTI_Base + mpc = MPC_LTI_Base.from_config(config, backend=bk) assert mpc.H is not None assert mpc.F is not None @@ -395,30 +418,35 @@ class TestSMC: def test_smc_construction_hurwitz_rejection(self, bk): """Non-Hurwitz surface coefficients raise ValueError.""" from shinro.controllers.smc import SlidingModeController + with pytest.raises(ValueError, match="Hurwitz"): SlidingModeController(c=[-1.0, 1.0], k1=1.0, backend=bk) def test_smc_construction_unknown_smoother(self, bk): """Unknown smoother name raises ValueError.""" from shinro.controllers.smc import SlidingModeController + with pytest.raises(ValueError, match="Unknown smoother"): SlidingModeController(c=[1.0, 2.0], k1=1.0, smoother="foo", backend=bk) def test_smc_n_property(self, bk): """n returns the length of the surface coefficient vector.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, backend=bk) assert ctrl.n == 2 def test_smc_hurwitz_polynomial_correct(self, bk): """c=[1,2] gives polynomial 2λ+1=0 with root at -0.5 (Hurwitz).""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, backend=bk) assert ctrl._is_hurwitz() def test_smc_hurwitz_polynomial_rejects_positive_root(self, bk): """c=[-1,1] gives polynomial λ-1=0 with root at +1 (not Hurwitz).""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, backend=bk) ctrl.c = ctrl.bk.array([-1.0, 1.0]) assert not ctrl._is_hurwitz() @@ -426,6 +454,7 @@ def test_smc_hurwitz_polynomial_rejects_positive_root(self, bk): def test_smc_compute_shape_scalar(self, bk): """compute() returns a 1-element array for a scalar-input system.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -436,6 +465,7 @@ def test_smc_compute_shape_scalar(self, bk): def test_smc_compute_shape_multi_input(self, bk): """compute() returns an m-element array for an m-input system.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 1.0], k1=1.0, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -443,9 +473,78 @@ def test_smc_compute_shape_multi_input(self, bk): u = ctrl.compute(x, f_x, g_x) assert _to_np(u, bk).shape == (2,) + def test_smc_multi_input_matches_min_norm_reference(self, bk): + """For n_u > 1 the closed form equals lstsq's minimum-norm solution. + + ``(c^T g) u = num`` is underdetermined (one surface, two inputs), so + the branch must return the minimum-norm member of the solution set — + the point ``np.linalg.lstsq`` picks, computed as ``cg^T (cg cg^T)^-1 + num`` so it can be lowered. + """ + from shinro.controllers.smc import SlidingModeController + + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.5, k2=0.5, phi=0.1, backend=bk) + rng = np.random.default_rng(0) + for _ in range(5): + x = bk.array(rng.normal(0.0, 0.5, 2)) + f_x = bk.array(rng.normal(0.0, 0.5, 2)) + g_x = bk.array(rng.normal(0.0, 0.5, (2, 2))) + u = _to_np(ctrl.compute(x, f_x, g_x), bk).ravel() + + c = _to_np(ctrl.c, bk) + s = float(c @ _to_np(x, bk)) + cf = float(c @ _to_np(f_x, bk)) + cg = c @ _to_np(g_x, bk) + smooth_s = np.clip(s / ctrl.phi, -1.0, 1.0) + s_dot_desired = -ctrl.k1 * abs(s) ** ctrl.alpha * smooth_s - ctrl.k2 * s + num = s_dot_desired - cf + want = np.linalg.lstsq(cg.reshape(1, -1), np.array([num]), rcond=None)[0] + + assert u.shape == (2,) + assert np.allclose(u, want, atol=1e-10) + # The chosen u delivers the desired reaching law along the surface: + # s_dot = c^T f + c^T g u == s_dot_desired. + assert np.allclose(cf + cg @ u, s_dot_desired, atol=1e-10) + + def test_smc_multi_input_min_norm_is_shortest(self, bk): + """The multi-input command is no longer than any other exact solution.""" + from shinro.controllers.smc import SlidingModeController + + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, phi=0.1, backend=bk) + x = bk.array([1.0, -0.5]) + f_x = bk.array([0.0, 0.0]) + g_x = bk.array([[1.0, 1.0], [0.0, 1.0]]) + u = _to_np(ctrl.compute(x, f_x, g_x), bk).ravel() + + c = _to_np(ctrl.c, bk) + cg = c @ _to_np(g_x, bk) + s = float(c @ _to_np(x, bk)) + num = -ctrl.k1 * abs(s) ** ctrl.alpha * np.clip(s / ctrl.phi, -1.0, 1.0) + # Any other exact solution is u + v with cg @ v == 0 and is strictly + # longer (or equal when v == 0). [cg[1], -cg[0]] spans the null space + # of the single row cg. + null_dir = np.array([cg[1], -cg[0]]) + for scale in (1.0, 0.5, -2.0): + v = scale * null_dir + assert np.isclose(cg @ v, 0.0, atol=1e-10) + assert np.linalg.norm(u + v) >= np.linalg.norm(u) - 1e-12 + assert np.isclose(cg @ u, num, atol=1e-10) + + def test_smc_multi_input_loss_of_controllability(self, bk): + """A near-zero ||c^T g|| raises RuntimeError on the multi-input branch.""" + from shinro.controllers.smc import SlidingModeController + + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, backend=bk) + x = bk.array([1.0, -0.5]) + f_x = bk.array([0.0, 0.0]) + g_x = bk.array([[0.0, 0.0], [0.0, 0.0]]) + with pytest.raises(RuntimeError, match="loss of controllability"): + ctrl.compute(x, f_x, g_x) + def test_smc_equivalent_control_analytical(self, bk): """For x_dot = f + g u with f=0, g=[0,1]^T, the equivalent control is u = -(c^T g)^{-1} c^T f = 0.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=0.0, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -459,6 +558,7 @@ def test_smc_equivalent_control_analytical(self, bk): def test_smc_equivalent_control_nonzero_f(self, bk): """For x_dot = f + g u with f=[0,1]^T, g=[0,1]^T, the control cancels f.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=0.0, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 1.0]) @@ -475,6 +575,7 @@ def test_smc_equivalent_control_nonzero_f(self, bk): def test_smc_sliding_surface_derivative_matches_desired(self, bk): """The actual s_dot = c^T f + c^T g u matches the desired reaching law.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.5, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -492,6 +593,7 @@ def test_smc_sliding_surface_derivative_matches_desired(self, bk): def test_smc_reaching_law_includes_k2_term(self, bk): """The reaching law includes the -k2*s term when k2 > 0.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, k2=3.0, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -509,6 +611,7 @@ def test_smc_reaching_law_includes_k2_term(self, bk): def test_smc_alpha_zero_gives_sign_law(self, bk): """With alpha=0, the reaching law is s_dot = -k1 * smooth(s) (|s|^0 = 1).""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=2.0, alpha=0.0, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -526,6 +629,7 @@ def test_smc_alpha_zero_gives_sign_law(self, bk): def test_smc_alpha_half_gives_sqrt_law(self, bk): """With alpha=0.5, the reaching law is s_dot = -k1 * |s|^0.5 * smooth(s).""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=2.0, alpha=0.5, phi=0.1, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -543,6 +647,7 @@ def test_smc_alpha_half_gives_sqrt_law(self, bk): def test_smc_sliding_surface_converges(self, bk): """The sliding surface s = c^T x converges toward zero under the control law.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=2.0, phi=0.05, backend=bk) dt = 0.01 x = bk.array([1.0, 0.0]) @@ -558,6 +663,7 @@ def test_smc_sliding_surface_converges(self, bk): def test_smc_sign_smoother_no_phi(self, bk): """With phi=0, the controller uses sign() and still drives s toward zero.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=2.0, phi=0.0, backend=bk) dt = 0.001 x = bk.array([1.0, 0.0]) @@ -573,6 +679,7 @@ def test_smc_sign_smoother_no_phi(self, bk): def test_smc_tanh_smoother(self, bk): """The tanh smoother produces a valid control action.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, phi=0.1, smoother="tanh", backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -583,6 +690,7 @@ def test_smc_tanh_smoother(self, bk): def test_smc_sigmoid_smoother(self, bk): """The sigmoid smoother produces a valid control action.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, phi=0.1, smoother="sigmoid", backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -593,6 +701,7 @@ def test_smc_sigmoid_smoother(self, bk): def test_smc_alpha_affects_convergence(self, bk): """Non-zero alpha changes the reaching law.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, phi=0.1, alpha=0.5, backend=bk) x = bk.array([1.0, -0.5]) f_x = bk.array([0.0, 0.0]) @@ -603,6 +712,7 @@ def test_smc_alpha_affects_convergence(self, bk): def test_smc_loss_of_controllability(self, bk): """A near-zero c^T g(x) raises RuntimeError.""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0, 0.0], k1=1.0, backend=bk) x = bk.array([1.0, 0.0, 0.0]) f_x = bk.array([0.0, 0.0, 0.0]) @@ -613,6 +723,7 @@ def test_smc_loss_of_controllability(self, bk): def test_smc_reset(self, bk): """reset() is a no-op (does not raise).""" from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController(c=[1.0, 2.0], k1=1.0, backend=bk) ctrl.reset() @@ -620,6 +731,7 @@ def test_smc_from_config(self, bk): """from_config creates a valid SMC controller.""" config = {"c": [1.0, 2.0], "k1": 1.0, "phi": 0.1} from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController.from_config(config, backend=bk) assert ctrl.n == 2 assert ctrl.k1 == 1.0 @@ -636,6 +748,7 @@ def test_smc_from_config_full(self, bk): "alpha": 0.3, } from shinro.controllers.smc import SlidingModeController + ctrl = SlidingModeController.from_config(config, backend=bk) assert ctrl.n == 3 assert ctrl.k1 == 2.0 @@ -651,6 +764,7 @@ class TestMPPI: def _ctrl(self, bk, **kwargs): """Build a minimal MPPI controller with identity dynamics and quadratic cost.""" from shinro.controllers.mppi import MPPIController + params = dict( dynamics_fn=lambda x, u, dt: bk.copy(x), cost_fn=lambda x, u: bk.sum(u**2, axis=1), @@ -668,6 +782,7 @@ def _ctrl(self, bk, **kwargs): def test_mppi_construction_validation(self, bk): """Invalid constructor parameters raise ValueError.""" from shinro.controllers.mppi import MPPIController + with pytest.raises(ValueError, match="num_samples"): MPPIController(num_samples=0, backend=bk) with pytest.raises(ValueError, match="temperature"): @@ -686,9 +801,14 @@ def test_mppi_compute_shape(self, bk): def test_mppi_compute_requires_callables(self, bk): """compute() without injected dynamics/cost raises RuntimeError.""" from shinro.controllers.mppi import MPPIController + ctrl = MPPIController( - num_samples=5, temperature=1.0, dt=0.1, horizon=3, - noise_sigma=[0.5], backend=bk, + num_samples=5, + temperature=1.0, + dt=0.1, + horizon=3, + noise_sigma=[0.5], + backend=bk, ) with pytest.raises(RuntimeError, match="dynamics_fn and cost_fn"): ctrl.compute(bk.array([0.0])) @@ -773,6 +893,30 @@ def test_mppi_nominal_sequence_shift(self, bk): expected_shift = np.concatenate([updated[1:], updated[-1:]]) assert np.allclose(ctrl.u, expected_shift, atol=1e-10) + def test_mppi_epsilon_param_matches_sampled(self, bk): + """Feeding epsilon reproduces the internally-sampled run exactly. + + This is the trace contract: the lowered kernel receives the + perturbations as an input port instead of sampling, so + ``compute(..., epsilon=...)`` must produce the same action the sampled + path did. Two controllers with different seeds must agree when one is + fed the other's draw. + """ + N, K = 6, 3 + sampled = self._ctrl(bk, num_samples=N, horizon=K, noise_sigma=[0.5], seed=5) + u_sampled = sampled.compute(bk.array([1.0, 1.0])) + eps = sampled._last_epsilon + costs_sampled = sampled._last_costs + assert eps is not None + assert costs_sampled is not None + assert eps.shape == (N, K, 1) + + fed = self._ctrl(bk, num_samples=N, horizon=K, noise_sigma=[0.5], seed=999) + u_fed = fed.compute(bk.array([1.0, 1.0]), None, np.ascontiguousarray(eps.reshape(N, K))) + assert fed._last_costs is not None + assert np.allclose(_to_np(u_sampled, bk), _to_np(u_fed, bk), atol=1e-12) + assert np.allclose(fed._last_costs, costs_sampled, atol=1e-12) + def test_mppi_bounds_clamp_in_rollout(self, bk): """The cost function never observes a control outside the configured bounds.""" bound = 0.3 @@ -837,6 +981,7 @@ def test_mppi_from_config(self, bk): "seed": 11, } from shinro.controllers.mppi import MPPIController + ctrl = MPPIController.from_config(config, backend=bk) assert ctrl.N == 50 assert ctrl.K == 8 @@ -847,6 +992,7 @@ def test_mppi_from_config(self, bk): def test_mppi_seed_reproducibility(self, bk): """The same seed produces identical control actions across instances.""" + def make(): return self._ctrl(bk, seed=42) @@ -872,10 +1018,16 @@ def test_mppi_attach_plant_lti(self, bk): """attach_plant wires an LTI plant's dynamics/cost into the controller.""" from shinro.controllers.mppi import MPPIController from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + plant = HolonomicMobileRobot(num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=0.02, backend=bk) ctrl = MPPIController( - num_samples=8, temperature=1.0, dt=0.02, horizon=4, - noise_sigma=[0.5, 0.5, 0.5], seed=1, backend=bk, + num_samples=8, + temperature=1.0, + dt=0.02, + horizon=4, + noise_sigma=[0.5, 0.5, 0.5], + seed=1, + backend=bk, ) ctrl.attach_plant(plant) u = ctrl.compute(bk.array([1.0, 2.0, 0.0])) @@ -885,10 +1037,16 @@ def test_mppi_attach_plant_nonlinear(self, bk): """attach_plant wires a nonlinear plant's dynamics/cost into the controller.""" from shinro.controllers.mppi import MPPIController from shinro.plants.inverted_pendulum import InvertedPendulum + plant = InvertedPendulum(backend=bk) ctrl = MPPIController( - num_samples=8, temperature=1.0, dt=0.01, horizon=4, - noise_sigma=[0.5], seed=1, backend=bk, + num_samples=8, + temperature=1.0, + dt=0.01, + horizon=4, + noise_sigma=[0.5], + seed=1, + backend=bk, ) ctrl.attach_plant(plant) u = ctrl.compute(bk.array([0.1, 0.0])) @@ -898,10 +1056,16 @@ def test_mppi_attach_plant_dimension_mismatch(self, bk): """attach_plant raises when the plant control dim disagrees with noise_sigma.""" from shinro.controllers.mppi import MPPIController from shinro.plants.inverted_pendulum import InvertedPendulum + plant = InvertedPendulum(backend=bk) ctrl = MPPIController( - num_samples=8, temperature=1.0, dt=0.01, horizon=4, - noise_sigma=[0.5, 0.5], seed=1, backend=bk, + num_samples=8, + temperature=1.0, + dt=0.01, + horizon=4, + noise_sigma=[0.5, 0.5], + seed=1, + backend=bk, ) with pytest.raises(ValueError, match="control dimension"): ctrl.attach_plant(plant) @@ -910,12 +1074,18 @@ def test_mppi_tracking_with_x_ref(self, bk): """With x_ref set, MPPI drives a plant toward the reference.""" from shinro.controllers.mppi import MPPIController from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + plant = HolonomicMobileRobot(num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=0.02, backend=bk) Q = bk.array([10.0, 10.0, 10.0]) R = bk.array([0.1, 0.1, 0.1]) ctrl = MPPIController( - num_samples=200, temperature=1.0, dt=0.02, horizon=10, - noise_sigma=[1.0, 1.0, 1.0], seed=1, backend=bk, + num_samples=200, + temperature=1.0, + dt=0.02, + horizon=10, + noise_sigma=[1.0, 1.0, 1.0], + seed=1, + backend=bk, ) ctrl.attach_plant(plant, Q=Q, R=R) x_ref = bk.array([1.0, 0.0, 0.0]) @@ -933,10 +1103,16 @@ def test_mppi_torch_backend_batched_ops(self, bk): pytest.skip("requires TorchBackend") from shinro.controllers.mppi import MPPIController from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + plant = HolonomicMobileRobot(num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=0.02, backend=bk) ctrl = MPPIController( - num_samples=8, temperature=1.0, dt=0.02, horizon=4, - noise_sigma=[0.5, 0.5, 0.5], seed=1, backend=bk, + num_samples=8, + temperature=1.0, + dt=0.02, + horizon=4, + noise_sigma=[0.5, 0.5, 0.5], + seed=1, + backend=bk, ) ctrl.attach_plant(plant) u = ctrl.compute(bk.array([1.0, 2.0, 0.0])) diff --git a/tests/test_measure.py b/tests/test_measure.py new file mode 100644 index 0000000..88e81a8 --- /dev/null +++ b/tests/test_measure.py @@ -0,0 +1,118 @@ +"""The kernel-size metric (``shinro.codegen.measure``). + +Pins the properties the metric relies on: the byte fields are derived from the +same manifest the VM compiles, they are self-consistent, and the MPPI graph it +measures is the standard one (``epsilon`` port ``(N, K*D_u)``, recurrent +``state_u``). The artifact measurement is exercised only when ``zig`` is +available — the metric is deliberately not part of the default CI gate. +""" + +import json +import shutil + +import pytest + +from shinro.codegen import measure +from shinro.codegen.lower_zig import lower_zig + + +class TestParseDims: + """The ``D_x x D_u x N x K`` spec parser.""" + + def test_valid(self): + assert measure.parse_dims("3x3x100x15") == (3, 3, 100, 15) + assert measure.parse_dims("6X6X10X4") == (6, 6, 10, 4) # case-insensitive + + def test_rejects_wrong_arity(self): + with pytest.raises(ValueError, match="D_x x D_u x N x K"): + measure.parse_dims("3x3x100") + + def test_rejects_non_numeric(self): + with pytest.raises(ValueError, match="four positive integers"): + measure.parse_dims("3x3xax15") + + def test_rejects_non_positive(self): + with pytest.raises(ValueError, match="positive"): + measure.parse_dims("3x3x0x15") + + +class TestGraphMetrics: + """Static size derivation from a lowered graph.""" + + def test_byte_fields_are_self_consistent(self, tmp_path): + cg = measure.mppi_lti_graph(3, 3, 6, 3) + m = measure.graph_metrics(cg, tmp_path) + + assert m["buf_bytes"] == m["buf_len_f64"] * 8 + assert m["input_bytes"] == sum(p["bytes"] for p in m["inputs"]) + assert m["output_bytes"] == sum(p["bytes"] for p in m["outputs"]) + assert m["state_bytes"] == sum(p["bytes"] for p in m["state_outputs"]) + assert m["nodes_total"] > 0 + assert m["has_solve_qp"] is False + + def test_manifest_records_the_memory_metric(self, tmp_path): + """The byte figures live in the manifest, so every graph self-reports.""" + lower_zig(measure.mppi_lti_graph(3, 3, 6, 3), str(tmp_path / "graph_data.zig")) + manifest = json.loads((tmp_path / "graph_data_manifest.json").read_text()) + + assert manifest["buf_bytes"] == manifest["buf_len"] * 8 + assert manifest["const_blob_bytes"] == manifest["const_blob_len"] * 8 + assert manifest["clip_blob_bytes"] == 2 * manifest["clip_blob_len"] * 8 + assert manifest["clip_blob_len"] > 0 # MPPI clips the sample batch + assert manifest["input_bytes"] == sum(p["bytes"] for p in manifest["inputs"]) + assert manifest["output_bytes"] == sum(p["bytes"] for p in manifest["outputs"]) + assert manifest["state_bytes"] == sum(p["bytes"] for p in manifest["state_outputs"]) + ports = manifest["inputs"] + manifest["outputs"] + manifest["state_outputs"] + assert all("bytes" in p for p in ports) + + def test_graph_metrics_reads_the_manifest(self, tmp_path): + """The metric is a faithful read — no duplicated buffer math.""" + m = measure.graph_metrics(measure.mppi_lti_graph(3, 3, 6, 3), tmp_path) + manifest = json.loads((tmp_path / "graph_data_manifest.json").read_text()) + for key in ("buf_bytes", "const_blob_bytes", "clip_blob_bytes", "input_bytes", "output_bytes", "state_bytes"): + assert m[key] == manifest[key], key + + def test_epsilon_dominates_the_input_buffer(self, tmp_path): + """The sampling port is N*K*D_u f64 — the host-side packing cost.""" + cg = measure.mppi_lti_graph(3, 3, 6, 3) + m = measure.graph_metrics(cg, tmp_path) + eps = next(p for p in m["inputs"] if p["name"] == "epsilon") + assert eps["shape"] == [6, 9] + assert eps["bytes"] == 6 * 3 * 3 * 8 + assert eps["bytes"] > m["input_bytes"] / 2 + + def test_state_port_is_the_nominal_plan(self, tmp_path): + cg = measure.mppi_lti_graph(4, 2, 8, 4) + m = measure.graph_metrics(cg, tmp_path) + assert [p["shape"] for p in m["state_outputs"]] == [[4, 2]] + + +class TestSweep: + """The multi-input sweep as a measurement (no compiler).""" + + def test_measure_static_records_every_config(self): + doc = measure.measure([(3, 3, 6, 3), (6, 6, 10, 4), (8, 4, 12, 5)], "ReleaseFast", build=False) + assert [c["label"] for c in doc["configs"]] == [ + "D_x=3 D_u=3 N=6 K=3", + "D_x=6 D_u=6 N=10 K=4", + "D_x=8 D_u=4 N=12 K=5", + ] + assert doc["build"] is False + # static-only: no artifact measured, but the size metrics are present + assert all(c["kernel"] is None for c in doc["configs"]) + assert all(c["graph"]["buf_bytes"] > 0 for c in doc["configs"]) + # wider inputs grow the input buffer (D_u=6 > D_u=3 at the same N) + by_label = {c["label"]: c["graph"]["input_bytes"] for c in doc["configs"]} + assert by_label["D_x=6 D_u=6 N=10 K=4"] > by_label["D_x=3 D_u=3 N=6 K=3"] + + +@pytest.mark.skipif(shutil.which("zig") is None, reason="zig not on PATH") +class TestKernelMetrics: + """The compiled-artifact measurement (opt-in: needs zig).""" + + def test_records_artifact_bytes_and_compile_cost(self, tmp_path): + cg = measure.mppi_lti_graph(3, 3, 6, 3) + k = measure.kernel_metrics(cg, tmp_path, "ReleaseFast") + assert k["optimize"] == "ReleaseFast" + assert k["so_bytes"] > 0 + assert k["compile_seconds"] > 0 diff --git a/tests/test_op_shape_matrix.py b/tests/test_op_shape_matrix.py index 1b9cb3a..c8d64cb 100644 --- a/tests/test_op_shape_matrix.py +++ b/tests/test_op_shape_matrix.py @@ -89,13 +89,22 @@ def _linalg_graph(g: Graph): outs["reshape_2d_to_1d"] = g.emit("reshape", [r23], (6,), target_shape=(6,)) specs = { - "a22": ((2, 2), "free"), "b22": ((2, 2), "free"), - "a21": ((2, 1), "free"), "s11": ((1, 1), "free"), - "v2": ((2,), "free"), "r12": ((1, 2), "free"), - "a31": ((3, 1), "free"), "b14": ((1, 4), "free"), - "t2": ((2, 2), "free"), "t23": ((2, 3), "free"), "t32": ((3, 2), "free"), - "i1": ((1, 1), "spd"), "i2": ((2, 2), "spd"), "i4": ((4, 4), "spd"), - "r6": ((6,), "free"), "r23": ((2, 3), "free"), + "a22": ((2, 2), "free"), + "b22": ((2, 2), "free"), + "a21": ((2, 1), "free"), + "s11": ((1, 1), "free"), + "v2": ((2,), "free"), + "r12": ((1, 2), "free"), + "a31": ((3, 1), "free"), + "b14": ((1, 4), "free"), + "t2": ((2, 2), "free"), + "t23": ((2, 3), "free"), + "t32": ((3, 2), "free"), + "i1": ((1, 1), "spd"), + "i2": ((2, 2), "spd"), + "i4": ((4, 4), "spd"), + "r6": ((6,), "free"), + "r23": ((2, 3), "free"), } return outs, specs @@ -122,7 +131,7 @@ def _elementwise_graph(g: Graph): v4b = g.input("v4b", (4, 1)) specs["v4b"] = ((4, 1), "free") - for op, oname in (("add", "add"), ("sub", "sub"), ("mul", "mul"), ("div", "div"), ("ne", "ne")): + for op, oname in (("add", "add"), ("sub", "sub"), ("mul", "mul"), ("div", "div"), ("ne", "ne"), ("lt", "lt")): outs[f"{oname}_same"] = g.emit(op, [x, same], (3, 2)) outs[f"{oname}_scalar"] = g.emit(op, [x, scalar], (3, 2)) outs[f"{oname}_row"] = g.emit(op, [x, row], (3, 2)) @@ -144,6 +153,12 @@ def _elementwise_graph(g: Graph): outs["clip_array"] = g.emit("clip", [c4], (4,), lo=np.full(4, -0.5), hi=np.full(4, 0.5)) outs["clip_scalar"] = g.emit("clip", [c4], (4,), lo=np.float64(-0.5), hi=np.float64(0.5)) outs["clip_2d_scalar"] = g.emit("clip", [c23], (2, 3), lo=np.float64(-0.4), hi=np.float64(0.4)) + # Per-channel bounds against a batched operand: numpy broadcasts (D,) over + # (N, D) — the shape MPPI clips its (N, D_u) sample batch with. The VM's + # clip blob needs one entry per element, expanded by lower_zig. + outs["clip_vec_broadcast"] = g.emit( + "clip", [c23], (2, 3), lo=np.array([-0.4, -0.2, -0.6]), hi=np.array([0.4, 0.2, 0.6]) + ) return outs, specs @@ -182,23 +197,37 @@ def _selection_graph(g: Graph): outs["any_2d"] = g.emit("any", [a23], ()) specs = { - "x32": ((3, 2), "free"), "x3": ((3,), "free"), - "s6": ((6,), "free"), "s42": ((4, 2), "free"), - "stack_a": ((2,), "free"), "stack_b": ((2,), "free"), "stack_c": ((2,), "free"), - "cp32": ((3, 2), "free"), "any4": ((4,), "free"), "any23": ((2, 3), "free"), + "x32": ((3, 2), "free"), + "x3": ((3,), "free"), + "s6": ((6,), "free"), + "s42": ((4, 2), "free"), + "stack_a": ((2,), "free"), + "stack_b": ((2,), "free"), + "stack_c": ((2,), "free"), + "cp32": ((3, 2), "free"), + "any4": ((4,), "free"), + "any23": ((2, 3), "free"), } return outs, specs def _pointwise_graph(g: Graph): - """tanh/relu/exp/sin/cos, argmax, one_hot.""" + """tanh/relu/exp/sin/cos/abs/sign, pow, argmax, one_hot.""" outs = {} x1 = g.input("p1", (1,)) x8 = g.input("p8", (8,)) - for op in ("tanh", "relu", "exp", "sin", "cos"): + for op in ("tanh", "relu", "exp", "sin", "cos", "abs", "sign"): outs[f"{op}_1"] = g.emit(op, [x1], (1,)) outs[f"{op}_8"] = g.emit(op, [x8], (8,)) + # pow on a positive base with a fractional exponent — the shape the SMC + # switching term uses (|s|^alpha). Free feeds can be negative, and + # pow(negative, fractional) is NaN on both engines; keeping the base abs'd + # keeps the cell deterministic rather than NaN-vs-NaN. + alpha = g.emit("const", [], (), value=np.float64(0.5)) + outs["pow_pos_frac_1"] = g.emit("pow", [g.emit("abs", [x1], (1,)), alpha], (1,)) + outs["pow_pos_frac_8"] = g.emit("pow", [g.emit("abs", [x8], (8,)), alpha], (8,)) + am5 = g.input("am5", (5,)) am23 = g.input("am23", (2, 3)) outs["argmax_1d"] = g.emit("argmax", [am5], ()) @@ -207,9 +236,21 @@ def _pointwise_graph(g: Graph): oh = g.input("oh_idx", (1,)) outs["one_hot_4"] = g.emit("one_hot", [oh], (4,), depth=4) + # min reduction cells: the full reduction (axis=None — MPPI's softmax + # shift), both 2-D axis reductions, and a NaN-propagation cell (0/0 via + # div) pinning numpy's NaN-wins semantics against the Zig comparison. + zero5 = g.emit("const", [], (5,), value=np.zeros(5)) + outs["min_1d"] = g.emit("min", [am5], (), axis=None) + outs["min_2d_all"] = g.emit("min", [am23], (), axis=None) + outs["min_2d_axis0"] = g.emit("min", [am23], (3,), axis=0) + outs["min_2d_axis1"] = g.emit("min", [am23], (2,), axis=1) + outs["min_nan_prop"] = g.emit("min", [g.emit("div", [am5, zero5], (5,))], (), axis=None) + specs = { - "p1": ((1,), "free"), "p8": ((8,), "free"), - "am5": ((5,), "free"), "am23": ((2, 3), "free"), + "p1": ((1,), "free"), + "p8": ((8,), "free"), + "am5": ((5,), "free"), + "am23": ((2, 3), "free"), "oh_idx": ((1,), "idx"), } return outs, specs @@ -316,9 +357,7 @@ def test_so_matches_numpy(self, matrix_so): for oname in cg.outputs: exp = np.asarray(traced[oname]) # shape contract: manifest declaration == numpy result - assert declared[oname] == exp.shape, ( - f"{name}/{oname}: manifest shape {declared[oname]} != numpy {exp.shape}" - ) + assert declared[oname] == exp.shape, f"{name}/{oname}: manifest shape {declared[oname]} != numpy {exp.shape}" got = out[off : off + exp.size] # NaN-aware: 0/0 in the boundary feeds must agree as NaN ok = np.isclose(got, exp.ravel(), rtol=0.0, atol=TOL, equal_nan=True) diff --git a/tests/test_plants.py b/tests/test_plants.py index 6e9caa5..e3dedba 100644 --- a/tests/test_plants.py +++ b/tests/test_plants.py @@ -636,3 +636,51 @@ def test_batch_mode(self, tmp_path): out_path.write_text(toml_string(config)) assert (output_dir / "pendulum.toml").exists() assert (output_dir / "cartpole.toml").exists() + + +class TestBatchCapableDynamics: + """``Plant.dynamics`` accepts a single state or a batch, consistently. + + Every nonlinear plant must be batch-capable — that is what lets the *same* + function serve the eager per-sample rollout, the finite-difference + linearization, and the lowered MPPI graph (see ``Plant.dynamics``). These + tests pin the rank contract and the batch/single agreement, so a new + nonlinear plant cannot silently ship with a scalar-only ``dynamics``. + """ + + def _plants(self, bk): + from shinro.plants.cartpole import CartPole + from shinro.plants.double_pendulum import DoublePendulum + from shinro.plants.inverted_pendulum import InvertedPendulum + return { + "inverted_pendulum": (InvertedPendulum(backend=bk), 2, 1), + "cartpole": (CartPole(backend=bk), 4, 1), + "double_pendulum": (DoublePendulum(backend=bk), 4, 2), + } + + def test_single_state_returns_single_derivative(self, bk): + rng = np.random.default_rng(0) + for name, (plant, n_x, n_u) in self._plants(bk).items(): + f = _to_np(plant.dynamics(bk.array(rng.normal(size=n_x)), bk.array(rng.normal(size=n_u))), bk) + assert f.shape == (n_x,), name + + def test_batch_matches_single_calls(self, bk): + """A ``(N, n_x)`` call equals ``N`` single-state calls, row for row.""" + rng = np.random.default_rng(1) + for name, (plant, n_x, n_u) in self._plants(bk).items(): + x_batch = bk.array(rng.normal(size=(5, n_x))) + u_batch = bk.array(rng.normal(size=(5, n_u))) + f_batch = _to_np(plant.dynamics(x_batch, u_batch), bk) + assert f_batch.shape == (5, n_x), name + for i in range(5): + f_i = _to_np(plant.dynamics(x_batch[i], u_batch[i]), bk) + assert np.allclose(f_batch[i], f_i, atol=1e-12), name + + def test_scalar_control_matches_vector(self, bk): + """A scalar control means the first input, the rest zero.""" + rng = np.random.default_rng(2) + for name, (plant, n_x, n_u) in self._plants(bk).items(): + x = bk.array(rng.normal(size=n_x)) + scalar = _to_np(plant.dynamics(x, 0.7), bk) + vector = _to_np(plant.dynamics(x, bk.array([0.7] + [0.0] * (n_u - 1))), bk) + assert np.allclose(scalar, vector, atol=1e-12), name diff --git a/tests/test_zig_lowering.py b/tests/test_zig_lowering.py index 48be2e8..6b9b204 100644 --- a/tests/test_zig_lowering.py +++ b/tests/test_zig_lowering.py @@ -32,7 +32,9 @@ from shinro.codegen.oracle import input_shape, output_split, pack_arrays, state_slices, step_so from shinro.codegen.trace_node import trace_node from shinro.codegen.tracing import Graph +from shinro.controllers.mppi import MPPIController from shinro.controllers.pid import PIDController +from shinro.controllers.smc import SlidingModeController from shinro.factories.controller_factory import ControllerFactory from shinro.factories.estimator_factory import EstimatorFactory from shinro.utils.array_backend import NumpyBackend @@ -135,8 +137,18 @@ def _build_lowered_ops_graph(): graph=g, inputs=["x"], outputs=[ - "tanh", "relu", "exp", "copy", "slice", "argmax", "one_hot", - "sin", "cos", "stack", "ne_zero", "ne_one", + "tanh", + "relu", + "exp", + "copy", + "slice", + "argmax", + "one_hot", + "sin", + "cos", + "stack", + "ne_zero", + "ne_one", ], state_inputs=[], state_outputs=[], @@ -152,9 +164,7 @@ def _build_mpc_graph(): (src/shinro/runtime/codegen/emosqp/), whose problem must match the ``mpc_lti_base.toml`` bake (n_vars=30). """ - ctrl = ControllerFactory( - str(REPO_ROOT / "src/shinro/configs/controllers/mpc_lti_base.toml") - ).create(backend=NumpyBackend()) + ctrl = ControllerFactory(str(REPO_ROOT / "src/shinro/configs/controllers/mpc_lti_base.toml")).create(backend=NumpyBackend()) ng = trace_node( ctrl, input_shapes={"current_state": (3,), "target_state": (3,)}, @@ -234,6 +244,228 @@ def _build_mpc_deltau_composed_graph(): return build_mpc_composed_graph("configs/controllers/mpc_base.toml") +def _smc_controller( + smoother: str = "sat", + phi: float = 0.1, + alpha: float = 0.0, + controllability_eps: float = 1e-12, +) -> SlidingModeController: + """The live (numpy) SMC the lowered graph is compared against. + + Same gains as the shipped ``configs/controllers/smc.toml`` (c=[1, 2], + k1=1, k2=0, sat, phi=0.1) so the fixture graph and the live component + cannot drift apart silently. + """ + return SlidingModeController( + c=[1.0, 2.0], + k1=1.0, + k2=0.0, + phi=phi, + smoother=smoother, + alpha=alpha, + controllability_eps=controllability_eps, + backend=NumpyBackend(), + ) + + +def _build_smc_graph( + smoother: str = "sat", + phi: float = 0.1, + alpha: float = 0.0, + controllability_eps: float = 1e-12, + n_u: int = 1, +) -> ComposedGraph: + """Standalone SMC graph: ``(x, f_x, g_x)`` in, ``(out, healthy)`` out. + + SMC is the first lowered controller whose runtime inputs are live plant + evaluations (``f_x = f(x)``, ``g_x = g(x)``), so there is no estimator to + compose with — and ``compose()`` deliberately has no role for ``f_x``/ + ``g_x`` (it refuses them rather than mis-wiring; see ``compose.py``). The + graph is therefore traced standalone and lowered directly, with the + dynamics terms as free C-ABI ports the host fills each tick. The plant + stays on the host, exactly as it does for every other lowered controller. + + Two outputs: ``out`` is the control (fail-safe-guarded — zero when the + controllability flag trips) and ``healthy`` is the flag SMC publishes via + :meth:`ArrayBackend.emit_named_output`. No recurrent state: SMC is + memoryless, so ``state_outputs`` is empty. + + Args: + smoother: ``sat`` / ``tanh`` / ``sigmoid`` boundary layer. + phi: Boundary layer thickness (0 selects the pure ``sign`` switch). + alpha: Fractional power on the switching term (exercises ``pow``). + controllability_eps: Near-zero ``‖c^T g‖`` threshold, baked into the + graph as a const — a plant-scaled deployment design parameter. + n_u: Number of control inputs. ``n_u == 1`` traces the division + branch; ``n_u > 1`` traces the minimum-norm pseudo-inverse branch + (``transpose`` + ``matmul``), with ``g_x`` shaped ``(2, n_u)``. + + Returns: + The traced, lowered-ready :class:`ComposedGraph`. + """ + smc = _smc_controller(smoother=smoother, phi=phi, alpha=alpha, controllability_eps=controllability_eps) + ng = trace_node(smc, input_shapes={"x": (2,), "f_x": (2,), "g_x": (2, n_u)}) + return ComposedGraph( + graph=ng.graph, + inputs=["x", "f_x", "g_x"], + outputs=["out", "healthy"], + state_inputs=[], + state_outputs=[], + ) + + +def _smc_rand_inputs(rng: np.random.Generator, min_cg: float = 0.2, n_u: int = 1) -> dict[str, np.ndarray]: + """Random SMC inputs with ``‖c^T g‖`` kept above the guard. + + ``c^T g = 0`` is measure-zero for continuous random data, so the samples + are rejection-filtered to stay clear of the fail-safe branch — this keeps + the numpy reference on its raising path, where it agrees with the graph. + For ``n_u == 1`` the norm check reduces to the old ``|c^T g|`` one. + """ + x = rng.normal(0.0, 0.5, (2,)) + f_x = rng.normal(0.0, 0.5, (2,)) + c = np.array([1.0, 2.0]) + while True: + g_x = rng.normal(0.0, 1.0, (2, n_u)) + if np.linalg.norm(c @ g_x) >= min_cg: + return {"x": x, "f_x": f_x, "g_x": g_x} + + +# MPPI graph dims stay deliberately small: the comptime VM unrolls the whole +# K-step rollout, so node count and the stack buffer grow with N*K*D_u. +MPPI_N, MPPI_K, MPPI_DX, MPPI_DU = 6, 3, 3, 3 + + +def _mppi_controller(N: int = MPPI_N, K: int = MPPI_K, dt: float = 0.02) -> MPPIController: + """The live (numpy) MPPI the lowered graph is compared against. + + Wired from an LTI plant through ``attach_plant`` — the production path + (``ScenarioFactory`` does the same) — so the traced graph and the live + reference cannot drift apart silently. The plant fixes D_x = D_u = 3. + """ + from shinro.plants.holonomicmobilerobot import HolonomicMobileRobot + + bk = NumpyBackend() + plant = HolonomicMobileRobot( + num_wheels=3, radius_robots=0.1, gamma=0.0, radius_wheels=0.03, dt=dt, backend=bk + ) + ctrl = MPPIController( + num_samples=N, + temperature=1.0, + dt=dt, + horizon=K, + noise_sigma=[0.5, 0.5, 0.5], + u_min=[-0.5, -0.5, -0.5], + u_max=[0.5, 0.5, 0.5], + seed=1, + backend=bk, + ) + ctrl.attach_plant(plant, Q=np.array([1.0, 1.0, 1.0]), R=np.array([0.1, 0.1, 0.1])) + return ctrl + + +def _build_mppi_graph(N: int = MPPI_N, K: int = MPPI_K, dt: float = 0.02) -> ComposedGraph: + """Standalone MPPI graph: the perturbations arrive as an input port. + + MPPI's Gaussian sampling stays on the host (see ``mppi.py``), so unlike + every other lowered controller there is no RNG in the graph: ``epsilon`` + is a free C-ABI port of shape ``(N, K*D_u)`` (sample-major) the host fills + each tick. That is what makes parity checkable — the same draw goes to the + graph and to the live controller. + + Recurrent state is the nominal control sequence ``u`` (``(K, D_u)``): the + controller rebinds it each tick, so ``trace_node`` detects it and the graph + emits ``state_u``. Two non-state outputs: ``out`` (the action) and + ``costs``, the per-sample rollout costs published through + :meth:`ArrayBackend.emit_named_output` — a host-visible diagnostic, the way + SMC publishes ``healthy``. + + Args: + N: Number of sampled perturbations. + K: Prediction horizon. + dt: Rollout time step (baked into the plant's model). + + Returns: + The traced, lowered-ready :class:`ComposedGraph`. + """ + ctrl = _mppi_controller(N=N, K=K, dt=dt) + ng = trace_node( + ctrl, + input_shapes={ + "current_state": (MPPI_DX,), + "target_state": (MPPI_DX,), + "epsilon": (N, K * MPPI_DU), + }, + state_shapes={"u": (K, MPPI_DU)}, + ) + return ComposedGraph( + graph=ng.graph, + inputs=["current_state", "target_state", "epsilon", "state_u"], + outputs=["out", "costs"], + state_inputs=["state_u"], + state_outputs=["state_u"], + ) + + +# The nonlinear rollout graph: MPPI on an InvertedPendulum (D_x = 2, D_u = 1). +MPPI_NL_DX, MPPI_NL_DU = 2, 1 + + +def _mppi_pendulum_controller(N: int = MPPI_N, K: int = MPPI_K, dt: float = 0.02) -> MPPIController: + """Live MPPI on a nonlinear plant — the reference the traced graph is checked against. + + Wired through ``attach_plant`` (the production path, as ``ScenarioFactory`` + does), so the rollout integrates the plant's own batch-capable + ``dynamics``. The plant fixes D_x = 2 (theta, theta_dot) and D_u = 1 (tau). + """ + from shinro.plants.inverted_pendulum import InvertedPendulum + + bk = NumpyBackend() + plant = InvertedPendulum(mass=0.1, length=0.5, damping=0.0, gravity=9.81, dt=dt, backend=bk) + ctrl = MPPIController( + num_samples=N, + temperature=1.0, + dt=dt, + horizon=K, + noise_sigma=[0.5], + u_min=[-0.5], + u_max=[0.5], + seed=1, + backend=bk, + ) + ctrl.attach_plant(plant, Q=np.array([1.0, 1.0]), R=np.array([0.1])) + return ctrl + + +def _build_mppi_pendulum_graph(N: int = MPPI_N, K: int = MPPI_K, dt: float = 0.02) -> ComposedGraph: + """Standalone nonlinear MPPI graph — same ports, a plant-dynamics rollout. + + The only difference from :func:`_build_mppi_graph` is the plant: the + rollout evaluates ``InvertedPendulum.dynamics`` over the whole sample batch + (``sin``/``mul`` nodes of shape ``(N, 1)``, not ``N`` copies of a scalar + body). The C-ABI contract is unchanged — ``epsilon`` in, ``out`` + ``costs`` + out, ``state_u`` recurrent — which is the point: the lowering path does not + care whether the dynamics are a matmul or a formula. + """ + ctrl = _mppi_pendulum_controller(N=N, K=K, dt=dt) + ng = trace_node( + ctrl, + input_shapes={ + "current_state": (MPPI_NL_DX,), + "target_state": (MPPI_NL_DX,), + "epsilon": (N, K * MPPI_NL_DU), + }, + state_shapes={"u": (K, MPPI_NL_DU)}, + ) + return ComposedGraph( + graph=ng.graph, + inputs=["current_state", "target_state", "epsilon", "state_u"], + outputs=["out", "costs"], + state_inputs=["state_u"], + state_outputs=["state_u"], + ) + + @pytest.fixture(scope="session") def base_so(tmp_path_factory): """Build the .so from the base_tracking composed graph once per session.""" @@ -299,6 +531,76 @@ def mpc_deltau_composed_so(tmp_path_factory, deltau_bake): ) +@pytest.fixture(scope="session") +def smc_so(tmp_path_factory): + """Build the .so for the shipped smc.toml config (sat, phi=0.1, alpha=0). + + Lowers to a tmp graph_path so the shared src/shinro/runtime/graph_data.zig + is not clobbered by this fixture (same discipline as the DeltaU case). + """ + d = tmp_path_factory.mktemp("zig-build-smc") + return _build_so(_build_smc_graph(), d, graph_path=d / "graph_data.zig") + + +@pytest.fixture(scope="session") +def smc_multi_so(tmp_path_factory): + """Build the .so for the n_u=2 SMC graph (minimum-norm pseudo-inverse branch). + + Two control inputs exercise the transpose/matmul closed form that the + scalar branch never emits. Tmp graph_path, same discipline as smc_so. + """ + d = tmp_path_factory.mktemp("zig-build-smc-multi") + return _build_so(_build_smc_graph(n_u=2), d, graph_path=d / "graph_data.zig") + + +@pytest.fixture(scope="session") +def mppi_so(tmp_path_factory): + """Build the .so from the standalone MPPI graph (the sampling-port contract). + + MPPI's perturbations arrive as a free C-ABI port, so this kernel does no + sampling: the host draws ``epsilon``, and three-way parity (.so vs + interpreter vs live numpy) is checkable exactly on that same draw. Lowers + to a tmp graph_path so the shared src/shinro/runtime/graph_data.zig is not + clobbered (same discipline as smc_so). + """ + d = tmp_path_factory.mktemp("zig-build-mppi") + return _build_so(_build_mppi_graph(), d, graph_path=d / "graph_data.zig") + + +@pytest.fixture(scope="session") +def mppi_pendulum_so(tmp_path_factory): + """Build the .so from the nonlinear MPPI graph (plant-``dynamics`` rollout). + + The same C-ABI contract as ``mppi_so`` — a free ``epsilon`` port, so parity + is checkable on one shared draw — but the rollout evaluates the plant's + batch-capable ``dynamics`` (``sin``/``mul`` nodes over the sample batch) + instead of a batched matmul. Lowers to a tmp graph_path so the shipped + runtime graph is never clobbered. + """ + d = tmp_path_factory.mktemp("zig-build-mppi-pendulum") + return _build_so(_build_mppi_pendulum_graph(), d, graph_path=d / "graph_data.zig") + + +# SMC config variants, each a graph-structure specialization: phi=0 swaps the +# clip boundary layer for the `sign` op, sigmoid adds the `abs` + `div` path, +# and alpha=0.5 exercises `pow` with a fractional exponent. +SMC_VARIANTS = [ + pytest.param({"phi": 0.0}, id="sign-phi0"), + pytest.param({"smoother": "sigmoid"}, id="sigmoid"), + pytest.param({"alpha": 0.5}, id="sat-alpha-pow"), +] + + +@pytest.fixture(scope="session", params=SMC_VARIANTS) +def smc_variant_so(request, tmp_path_factory): + """Build a .so per SMC config variant (each is its own lowered graph).""" + cg = _build_smc_graph(**request.param) + slug = "-".join(f"{k}-{v}" for k, v in sorted(request.param.items())) + d = tmp_path_factory.mktemp(f"zig-build-smc-{slug}") + lib, composed = _build_so(cg, d, graph_path=d / "graph_data.zig") + return lib, composed, request.param + + def _pack_inputs(cg, y, x_ref, u_prev, x_hat_init, P_init): """Pack host inputs into the flat C-ABI buffer, in cg.inputs order.""" port_arrays = { @@ -475,11 +777,7 @@ def test_so_matches_numpy_for_every_shape(self, matmul_shapes_so): offsets = {} off = 0 for name in cg.outputs: - size = next( - int(np.prod(n.shape)) - for n in cg.graph.nodes - if n.op == "output" and n.attrs["name"] == name - ) + size = next(int(np.prod(n.shape)) for n in cg.graph.nodes if n.op == "output" and n.attrs["name"] == name) offsets[name] = (off, off + size) off += size @@ -599,11 +897,7 @@ def test_so_matches_interpret_single_input(self, single_input_kf_lqr_so): #: Synthetic dimensionality sweep at the Quadrotor-scale: n_x x n_u combos #: beyond any named plant, LQR+KF. Exercises large matmuls and the n_x x n_x #: Kalman inverse (12x12, 24x24) the named plants never reach. -DIM_SWEEP_CASES = [ - (f"synth{nx}x{nu}", "synthetic", None, nx, nu, 0.01, "LQR", "KalmanFilter") - for nx in (6, 12, 24) - for nu in (1, 3, 4) -] +DIM_SWEEP_CASES = [(f"synth{nx}x{nu}", "synthetic", None, nx, nu, 0.01, "LQR", "KalmanFilter") for nx in (6, 12, 24) for nu in (1, 3, 4)] ALL_SCAN_CASES = PLANT_SCAN_CASES + DIM_SWEEP_CASES @@ -727,9 +1021,7 @@ def _scan_input_ports(cg, n_x, n_u, rng): if name in known: ports[name] = known[name] else: - shape = next( - n.shape for n in cg.graph.nodes if n.op == "input" and n.attrs["name"] == name - ) + shape = next(n.shape for n in cg.graph.nodes if n.op == "input" and n.attrs["name"] == name) ports[name] = np.zeros(tuple(shape)) return ports @@ -750,53 +1042,70 @@ def plant_so(tmp_path_factory, request): def _glue_probe_case(name): """(builder, in_specs, feed) for one glue-op shape probe.""" if name == "transpose-nonsquare": + def build(g): return {"t": g.emit("transpose", [g.input("x", (2, 3))], (3, 2))} + return build, [("x", (2, 3))], {"x": np.arange(6, dtype=float).reshape(2, 3)} if name == "clip-scalar-bounds": + def build(g): x = g.input("x", (4,)) return {"c": g.emit("clip", [x], (4,), lo=np.float64(-0.5), hi=np.float64(0.5))} + return build, [("x", (4,))], {"x": np.array([-2.0, -0.1, 0.1, 2.0])} if name == "where-scalar-branch": + def build(g): x = g.input("x", (3,)) one = g.emit("const", [], (), value=np.float64(1.0)) zero = g.emit("const", [], (), value=np.float64(0.0)) cond = g.emit("ne", [x, zero], (3,)) return {"w": g.emit("where", [cond, one, x], (3,))} + return build, [("x", (3,))], {"x": np.array([0.0, 1.0, 2.0])} if name == "where-row-broadcast": + def build(g): x = g.input("x", (3, 2)) bias = g.emit("const", [], (1, 2), value=np.array([[1.0, 2.0]])) zero = g.emit("const", [], (1, 2), value=np.zeros((1, 2))) cond = g.emit("ne", [x, zero], (3, 2)) return {"w": g.emit("where", [cond, bias, x], (3, 2))} + return build, [("x", (3, 2))], {"x": np.arange(6, dtype=float).reshape(3, 2)} if name == "ew2-row-broadcast": + def build(g): x = g.input("x", (3, 2)) bias = g.emit("const", [], (1, 2), value=np.array([[10.0, 20.0]])) return {"s": g.emit("add", [x, bias], (3, 2))} + return build, [("x", (3, 2))], {"x": np.ones((3, 2))} if name == "ew2-col-broadcast": + def build(g): x = g.input("x", (3, 2)) scale = g.emit("const", [], (3, 1), value=np.array([[2.0], [3.0], [4.0]])) return {"s": g.emit("mul", [x, scale], (3, 2))} + return build, [("x", (3, 2))], {"x": np.full((3, 2), 1.5)} if name == "slice-2d-rows": + def build(g): x = g.input("x", (4, 2)) return {"s": g.emit("slice", [x], (2, 2), start=1, stop=3)} + return build, [("x", (4, 2))], {"x": np.arange(8, dtype=float).reshape(4, 2)} if name == "slice-1d-control": + def build(g): x = g.input("x", (6,)) return {"s": g.emit("slice", [x], (3,), start=2, stop=5)} + return build, [("x", (6,))], {"x": np.arange(6, dtype=float)} if name == "transcendentals": + def build(g): x = g.input("x", (8,)) return { @@ -805,6 +1114,7 @@ def build(g): "sin": g.emit("sin", [x], (8,)), "cos": g.emit("cos", [x], (8,)), } + rng = np.random.default_rng(7) return build, [("x", (8,))], {"x": rng.normal(0, 2, 8)} raise ValueError(name) @@ -812,11 +1122,11 @@ def build(g): GLUE_CASES = [ # the five found-bug cells ... - "transpose-nonsquare", # VM had square-only stride symmetry -> silent garbage - "clip-scalar-bounds", # scalar bounds -> flat-blob comptime OOB - "where-scalar-branch", # scalar branch -> runtime OOB panic - "ew2-row-broadcast", # (1,2)+(3,2) -> runtime OOB panic - "slice-2d-rows", # flat-offset indexing on a row slice -> silent garbage + "transpose-nonsquare", # VM had square-only stride symmetry -> silent garbage + "clip-scalar-bounds", # scalar bounds -> flat-blob comptime OOB + "where-scalar-branch", # scalar branch -> runtime OOB panic + "ew2-row-broadcast", # (1,2)+(3,2) -> runtime OOB panic + "slice-2d-rows", # flat-offset indexing on a row slice -> silent garbage # ... and broadcast/shape cells adjacent to them "where-row-broadcast", "ew2-col-broadcast", @@ -912,9 +1222,7 @@ def test_compose_rejects_nonsquare_pid(tmp_path): ) ctrl = tmp_path / "pid_cartpole.toml" ctrl.write_text( - 'type = "PID"\nname = "pid"\ndt = 0.01\n' - "kp = [2.0]\nki = [0.5]\nkd = [0.5]\n" - "output_limits = { min = [-10.0], max = [10.0] }\n" + 'type = "PID"\nname = "pid"\ndt = 0.01\nkp = [2.0]\nki = [0.5]\nkd = [0.5]\noutput_limits = { min = [-10.0], max = [10.0] }\n' ) with _pytest.raises(ValueError, match="control dimension"): build_composed_graph(str(est), str(ctrl), n_x=4, n_u=1) @@ -967,9 +1275,513 @@ def test_lowered_ops_match_interpreter(self, lowered_ops_so): np.testing.assert_allclose(got, expected, rtol=1e-14, atol=1e-14) else: assert name in exact_ops, f"unexpected op {name}" - assert np.array_equal(got, expected), ( - f"op {name} diverged: got {got}, expected {expected}" - ) + assert np.array_equal(got, expected), f"op {name} diverged: got {got}, expected {expected}" + + +class TestSmcOracle: + """The lowered SMC control law matches the interpreter and live numpy. + + SMC is the first lowered controller whose runtime inputs are live plant + evaluations, so the graph is standalone (no estimator) with ``f_x``/ + ``g_x`` as free C-ABI ports. This suite covers the ops added for it + (``abs`` / ``sign`` / ``pow`` / ``lt``) plus the auxiliary ``healthy`` + port and the where-guarded fail-safe: a graph has no exceptions, so the + near-zero ``c^T g`` guard compiles to a zero command + a flag instead of + a raise. + """ + + def test_so_matches_interpreter_and_numpy(self, smc_so): + """.so, graph interpreter, and live numpy agree on 25 seeded samples.""" + lib, cg = smc_so + n_out, n_state = output_split(cg) + assert n_state == 0 + assert n_out == 2 # out + healthy + + rng = np.random.default_rng(11) + smc = _smc_controller() + max_err = 0.0 + for _ in range(25): + arrays = _smc_rand_inputs(rng) + out, _ = step_so(lib, pack_arrays(cg, arrays), n_out, n_state) + traced = interpret(cg.graph, arrays) + want = np.asarray(smc.compute(arrays["x"], arrays["f_x"], arrays["g_x"])).ravel() + + # healthy = 1 on this path (|c^T g| ≥ 0.2 ≫ 1e-12) + assert out[1] == 1.0 + assert traced["healthy"][0] == 1.0 + np.testing.assert_allclose(out[0], traced["out"][0], rtol=1e-14, atol=1e-14) + np.testing.assert_allclose(out[0], want[0], rtol=1e-12, atol=1e-12) + max_err = max(max_err, abs(out[0] - want[0])) + assert max_err < 1e-12, f"SMC .so drifted from live numpy: {max_err:.3e}" + + @pytest.mark.parametrize( + "g_x", + [ + np.array([[1.0], [-0.5]]), # c^T g == 0 exactly + np.array([[0.0], [0.0]]), # degenerate actuator channel + ], + ids=["exact-zero", "zero-g"], + ) + def test_lost_controllability_is_failsafe_and_flagged(self, smc_so, g_x): + """A graph cannot raise: |c^T g| below eps → u == 0 and healthy == 0. + + The compiled analogue of numpy's RuntimeError. Zero-command is the + kernel's floor, not a safety guarantee — the flag is what lets the + host run its own fault policy in the same tick. + """ + lib, cg = smc_so + n_out, n_state = output_split(cg) + arrays = {"x": np.array([1.0, 0.5]), "f_x": np.array([0.3, -0.2]), "g_x": g_x} + + out, _ = step_so(lib, pack_arrays(cg, arrays), n_out, n_state) + assert out[0] == 0.0, "fail-safe must emit exactly zero" + assert out[1] == 0.0, "controllability flag must be low" + + # The interpreter (the graph's own reference) agrees, and the live + # component still raises — the two paths differ only here, by design. + with np.errstate(divide="ignore", invalid="ignore"): + traced = interpret(cg.graph, arrays) + assert traced["out"][0] == 0.0 + assert traced["healthy"][0] == 0.0 + with pytest.raises(RuntimeError, match="near-zero"): + _smc_controller().compute(arrays["x"], arrays["f_x"], arrays["g_x"]) + + def test_so_matches_interpreter_and_numpy_multi_input(self, smc_multi_so): + """n_u=2: .so, interpreter, and live numpy agree on the min-norm branch.""" + lib, cg = smc_multi_so + n_out, n_state = output_split(cg) + assert n_state == 0 + assert n_out == 3 # out (2 inputs) + healthy + + rng = np.random.default_rng(17) + smc = _smc_controller() + c = np.array([1.0, 2.0]) + max_err = 0.0 + for _ in range(25): + arrays = _smc_rand_inputs(rng, n_u=2) + out, _ = step_so(lib, pack_arrays(cg, arrays), n_out, n_state) + traced = interpret(cg.graph, arrays) + want = np.asarray(smc.compute(arrays["x"], arrays["f_x"], arrays["g_x"])).ravel() + + # healthy = 1 on this path (||c^T g|| >= 0.2 >> 1e-12) + assert out[2] == 1.0 + assert traced["healthy"][0] == 1.0 + np.testing.assert_allclose(out[:2], traced["out"], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(out[:2], want, rtol=1e-12, atol=1e-12) + + # The command delivers the reaching law along the surface: the + # pseudo-inverse is not just close to numpy, it solves the system. + s = float(c @ arrays["x"]) + num = -smc.k1 * abs(s) ** smc.alpha * np.clip(s / smc.phi, -1.0, 1.0) - float(c @ arrays["f_x"]) + assert np.isclose((c @ arrays["g_x"]) @ want, num, atol=1e-10) + max_err = max(max_err, float(np.max(np.abs(out[:2] - want)))) + assert max_err < 1e-12, f"SMC n_u=2 .so drifted from live numpy: {max_err:.3e}" + + def test_multi_input_lost_controllability_is_failsafe_and_flagged(self, smc_multi_so): + """n_u>1: ||c^T g|| below eps → u == 0 (both inputs) and healthy == 0.""" + lib, cg = smc_multi_so + n_out, n_state = output_split(cg) + # c = [1, 2]; each column is orthogonal to c, so c^T g == [0, 0]. + g_x = np.array([[2.0, -2.0], [-1.0, 1.0]]) + assert np.allclose(np.array([1.0, 2.0]) @ g_x, 0.0) + arrays = {"x": np.array([1.0, 0.5]), "f_x": np.array([0.3, -0.2]), "g_x": g_x} + + out, _ = step_so(lib, pack_arrays(cg, arrays), n_out, n_state) + np.testing.assert_array_equal(out[:2], np.zeros(2)) + assert out[2] == 0.0, "controllability flag must be low" + + with np.errstate(divide="ignore", invalid="ignore"): + traced = interpret(cg.graph, arrays) + np.testing.assert_array_equal(traced["out"], np.zeros(2)) + assert traced["healthy"][0] == 0.0 + with pytest.raises(RuntimeError, match="near-zero"): + _smc_controller().compute(arrays["x"], arrays["f_x"], arrays["g_x"]) + + def test_multi_input_graph_uses_the_pseudo_inverse_ops(self, smc_multi_so): + """The n_u>1 branch is matmul + transpose, not a baked solve constant.""" + _, cg = smc_multi_so + ops = {node.op for node in cg.graph.nodes} + assert {"transpose", "matmul", "abs", "lt", "where"} <= ops, f"missing ops: {sorted(ops)}" + assert set(cg.outputs) == {"out", "healthy"} + assert cg.state_outputs == [] # SMC is memoryless + + def test_graph_uses_the_new_ops_and_aux_port(self, smc_so): + """Drift guard: the guard/flag structure is actually in the graph.""" + _, cg = smc_so + ops = {node.op for node in cg.graph.nodes} + assert {"abs", "pow", "lt", "where"} <= ops, f"missing guard ops: {sorted(ops)}" + assert set(cg.outputs) == {"out", "healthy"} + assert cg.state_outputs == [] # SMC is memoryless + + def test_graph_contains_sign_for_pure_switching(self, smc_variant_so): + """phi == 0 selects the sign path; the variant graph must contain it.""" + _, cg, cfg = smc_variant_so + if cfg.get("phi") != 0.0: + pytest.skip("sign only appears when phi == 0") + ops = {node.op for node in cg.graph.nodes} + assert "sign" in ops + + def test_variant_graphs_match_interpreter(self, smc_variant_so): + """Each config variant is its own graph and matches the interpreter. + + phi=0 swaps clip for sign; sigmoid (s/(|s|+phi)) adds abs; alpha=0.5 + exercises pow with a fractional exponent — all against the graph's + own interpreter reference, with the live numpy component as the + second opinion. + """ + lib, cg, cfg = smc_variant_so + n_out, n_state = output_split(cg) + rng = np.random.default_rng(29) + smc = _smc_controller(**cfg) + + for _ in range(15): + arrays = _smc_rand_inputs(rng) + out, _ = step_so(lib, pack_arrays(cg, arrays), n_out, n_state) + traced = interpret(cg.graph, arrays) + want = np.asarray(smc.compute(arrays["x"], arrays["f_x"], arrays["g_x"])).ravel() + assert out[1] == 1.0 + np.testing.assert_allclose(out[0], traced["out"][0], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(out[0], want[0], rtol=1e-12, atol=1e-12) + + def test_controllability_eps_is_baked_from_config(self, tmp_path): + """eps is a deployment knob: a plant-scaled value trips the flag earlier. + + With eps = 0.5, a perfectly usable-but-small |c^T g| = 0.2 is treated + as lost controllability — the arithmetic-only 1e-12 default would let + it through and amplify 1/0.2 instead. The graph is built with its own + graph_path so the shared src/shinro/runtime/graph_data.zig is untouched. + """ + cg = _build_smc_graph(controllability_eps=0.5) + lib, _ = _build_so(cg, tmp_path, graph_path=tmp_path / "graph_data.zig") + n_out, n_state = output_split(cg) + arrays = {"x": np.array([1.0, 0.5]), "f_x": np.array([0.3, -0.2]), "g_x": np.array([[0.2], [0.0]])} + assert abs(0.2) > 1e-12 # the default guard would not fire here + + out, _ = step_so(lib, pack_arrays(cg, arrays), n_out, n_state) + assert out[0] == 0.0 + assert out[1] == 0.0 + with np.errstate(divide="ignore", invalid="ignore"): + traced = interpret(cg.graph, arrays) + assert traced["healthy"][0] == 0.0 + + +class TestMppiOracle: + """The interpreted MPPI graph matches live numpy — including recurrence. + + MPPI's Gaussian sampling stays on the host, so the perturbations arrive + through a free ``epsilon`` port and the traced call never touches the RNG. + That is what makes parity checkable at all: feed the same draw to the graph + and to the live controller and they must agree, tick after tick. + + These cases are interpreter-only (no .so): they pin that the traced graph + computes the same control law the live component does, which is the oracle + the Zig VM is checked against next. + """ + + def _feeds(self, x0, x_ref, epsilon, state_u=None): + """The four C-ABI input ports; ``state_u`` defaults to a zero plan.""" + return { + "current_state": x0, + "target_state": x_ref, + "epsilon": epsilon, + "state_u": np.zeros((MPPI_K, MPPI_DU)) if state_u is None else state_u, + } + + def test_interpreter_matches_numpy(self): + """interpret() == live numpy across 10 seeded perturbation draws.""" + cg = _build_mppi_graph() + ref = _mppi_controller() + rng = np.random.default_rng(23) + x_ref = np.array([1.0, 0.0, 0.0]) + max_u_err = 0.0 + max_state_err = 0.0 + for _ in range(10): + # Each iteration is a fresh tick: the graph is fed a zero plan, so + # the live reference must start from one too. + ref.reset() + x0 = rng.normal(0.0, 0.5, MPPI_DX) + eps = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)) + out = interpret(cg.graph, self._feeds(x0, x_ref, eps)) + want_u = np.asarray(ref.compute(x0, x_ref, eps)).ravel() + want_state = np.asarray(ref.u).reshape(MPPI_K, MPPI_DU) + + np.testing.assert_allclose(out["out"], want_u, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose(out["state_u"], want_state, rtol=1e-11, atol=1e-11) + max_u_err = max(max_u_err, np.max(np.abs(out["out"] - want_u))) + max_state_err = max(max_state_err, np.max(np.abs(out["state_u"] - want_state))) + assert max_u_err < 1e-11, f"MPPI graph drifted from live numpy: {max_u_err:.3e}" + assert max_state_err < 1e-11 + + def test_costs_port_is_published(self): + """The per-sample rollout costs are a graph output port (diagnostic).""" + cg = _build_mppi_graph() + rng = np.random.default_rng(31) + out = interpret( + cg.graph, + self._feeds( + rng.normal(0.0, 0.5, MPPI_DX), + np.array([1.0, 0.0, 0.0]), + rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)), + ), + ) + assert out["costs"].shape == (MPPI_N,) + assert np.all(np.isfinite(out["costs"])) + + def test_recurrence_matches_sequential_ticks(self): + """Feeding state_u back reproduces a second live tick (recurrent edge).""" + cg = _build_mppi_graph() + ref = _mppi_controller() + rng = np.random.default_rng(29) + x0 = rng.normal(0.0, 0.5, MPPI_DX) + x_ref = np.array([1.0, 0.0, 0.0]) + eps1 = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)) + eps2 = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)) + + tick1 = interpret(cg.graph, self._feeds(x0, x_ref, eps1)) + want1 = np.asarray(ref.compute(x0, x_ref, eps1)).ravel() + np.testing.assert_allclose(tick1["out"], want1, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose( + tick1["state_u"], np.asarray(ref.u).reshape(MPPI_K, MPPI_DU), rtol=1e-11, atol=1e-11 + ) + + # The graph's own state output feeds the next tick — no numpy state. + tick2 = interpret(cg.graph, self._feeds(x0, x_ref, eps2, state_u=tick1["state_u"])) + want2 = np.asarray(ref.compute(x0, x_ref, eps2)).ravel() + np.testing.assert_allclose(tick2["out"], want2, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose( + tick2["state_u"], np.asarray(ref.u).reshape(MPPI_K, MPPI_DU), rtol=1e-11, atol=1e-11 + ) + + def test_graph_structure_and_ports(self): + """Drift guard: the ops MPPI relies on and the C-ABI port layout.""" + cg = _build_mppi_graph() + ops = {node.op for node in cg.graph.nodes} + for op in ("min", "matmul", "clip", "slice", "stack", "transpose", "exp", "reshape"): + assert op in ops, f"MPPI graph lost the {op!r} op" + + assert cg.outputs == ["out", "costs"] + assert cg.state_outputs == ["state_u"] + port_shapes = {n.attrs["name"]: n.shape for n in cg.graph.nodes if n.op == "input"} + # The sampling contract: (N, K*D_u), sample-major — what the host packs. + assert port_shapes["epsilon"] == (MPPI_N, MPPI_K * MPPI_DU) + assert port_shapes["state_u"] == (MPPI_K, MPPI_DU) + assert port_shapes["current_state"] == (MPPI_DX,) + assert port_shapes["target_state"] == (MPPI_DX,) + + def test_so_matches_interpreter_and_numpy(self, mppi_so): + """.so, interpreter, and live numpy agree on the same seeded draws.""" + lib, cg = mppi_so + n_out, n_state = output_split(cg) + assert n_state == MPPI_K * MPPI_DU # the nominal plan recurs + assert n_out == MPPI_DU + MPPI_N # out (D_u) + costs (N) + + ref = _mppi_controller() + rng = np.random.default_rng(41) + x_ref = np.array([1.0, 0.0, 0.0]) + max_u_err = 0.0 + for _ in range(10): + ref.reset() + x0 = rng.normal(0.0, 0.5, MPPI_DX) + eps = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)) + feeds = self._feeds(x0, x_ref, eps) + out, state = step_so(lib, pack_arrays(cg, feeds), n_out, n_state) + traced = interpret(cg.graph, feeds) + want_u = np.asarray(ref.compute(x0, x_ref, eps)).ravel() + + # kernel vs its own interpreter (the tight tier), then vs numpy + np.testing.assert_allclose(out[:MPPI_DU], traced["out"], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(state, np.asarray(traced["state_u"]).ravel(), rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(out[MPPI_DU:], np.asarray(traced["costs"]).ravel(), rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(out[:MPPI_DU], want_u, rtol=1e-11, atol=1e-11) + max_u_err = max(max_u_err, np.max(np.abs(out[:MPPI_DU] - want_u))) + assert max_u_err < 1e-11, f"MPPI .so drifted from live numpy: {max_u_err:.3e}" + + def test_cabi_recurrence_matches_numpy(self, mppi_so): + """The kernel's own state buffer reproduces a second live tick. + + The host feeds ``state_out`` straight back as the next tick's + ``state_u`` — no numpy state in the loop. Two ticks must match two + sequential live ``compute()`` calls. + """ + lib, cg = mppi_so + n_out, n_state = output_split(cg) + start, stop = state_slices(cg)["state_u"] + ref = _mppi_controller() + rng = np.random.default_rng(43) + x0 = rng.normal(0.0, 0.5, MPPI_DX) + x_ref = np.array([1.0, 0.0, 0.0]) + eps1 = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)) + eps2 = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_DU)) + + out1, state1 = step_so(lib, pack_arrays(cg, self._feeds(x0, x_ref, eps1)), n_out, n_state) + want1 = np.asarray(ref.compute(x0, x_ref, eps1)).ravel() + np.testing.assert_allclose(out1[:MPPI_DU], want1, rtol=1e-11, atol=1e-11) + + # Feed the kernel's state back through the C-ABI input buffer. + plan = state1[start:stop].reshape(MPPI_K, MPPI_DU) + out2, state2 = step_so(lib, pack_arrays(cg, self._feeds(x0, x_ref, eps2, state_u=plan)), n_out, n_state) + want2 = np.asarray(ref.compute(x0, x_ref, eps2)).ravel() + np.testing.assert_allclose(out2[:MPPI_DU], want2, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose(state2[start:stop], np.asarray(ref.u).ravel(), rtol=1e-11, atol=1e-11) + + +class TestMppiNonlinearOracle: + """The nonlinear MPPI graph: same C-ABI ports, a plant-``dynamics`` rollout. + + MPPI on a nonlinear plant rolls out with the plant's own batch-capable + ``dynamics``, so the graph contains one ``sin``/``mul`` node of shape + ``(N, 1)`` per formula term instead of ``N`` copies of a scalar body — the + property that keeps the node count independent of the sample count. These + cases are interpreter-only; the ``.so`` three-way parity follows. + """ + + def _feeds(self, x0, x_ref, epsilon, state_u=None): + """The four C-ABI input ports; ``state_u`` defaults to a zero plan.""" + return { + "current_state": x0, + "target_state": x_ref, + "epsilon": epsilon, + "state_u": np.zeros((MPPI_K, MPPI_NL_DU)) if state_u is None else state_u, + } + + def _draw(self, rng): + """A tick's state, reference, and perturbation draw.""" + x0 = rng.normal(0.0, 0.3, MPPI_NL_DX) + x_ref = np.array([0.5, 0.0]) + eps = rng.normal(0.0, 0.5, (MPPI_N, MPPI_K * MPPI_NL_DU)) + return x0, x_ref, eps + + def test_interpreter_matches_numpy(self): + """interpret() == live numpy across 10 seeded draws (nonlinear rollout).""" + cg = _build_mppi_pendulum_graph() + ref = _mppi_pendulum_controller() + rng = np.random.default_rng(23) + max_u_err = 0.0 + max_state_err = 0.0 + for _ in range(10): + ref.reset() + x0, x_ref, eps = self._draw(rng) + out = interpret(cg.graph, self._feeds(x0, x_ref, eps)) + want_u = np.asarray(ref.compute(x0, x_ref, eps)).ravel() + want_state = np.asarray(ref.u).reshape(MPPI_K, MPPI_NL_DU) + np.testing.assert_allclose(out["out"], want_u, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose(out["state_u"], want_state, rtol=1e-11, atol=1e-11) + max_u_err = max(max_u_err, float(np.max(np.abs(out["out"] - want_u)))) + max_state_err = max(max_state_err, float(np.max(np.abs(out["state_u"] - want_state)))) + assert max_u_err < 1e-11, f"nonlinear MPPI graph drifted from live numpy: {max_u_err:.3e}" + assert max_state_err < 1e-11 + + def test_costs_port_is_published(self): + """The per-sample rollout costs remain a graph output port.""" + cg = _build_mppi_pendulum_graph() + rng = np.random.default_rng(31) + out = interpret(cg.graph, self._feeds(*self._draw(rng))) + assert out["costs"].shape == (MPPI_N,) + assert np.all(np.isfinite(out["costs"])) + + def test_recurrence_matches_sequential_ticks(self): + """Feeding state_u back reproduces a second live tick.""" + cg = _build_mppi_pendulum_graph() + ref = _mppi_pendulum_controller() + rng = np.random.default_rng(29) + x0, x_ref, eps1 = self._draw(rng) + _, _, eps2 = self._draw(rng) + + tick1 = interpret(cg.graph, self._feeds(x0, x_ref, eps1)) + want1 = np.asarray(ref.compute(x0, x_ref, eps1)).ravel() + np.testing.assert_allclose(tick1["out"], want1, rtol=1e-11, atol=1e-11) + + tick2 = interpret(cg.graph, self._feeds(x0, x_ref, eps2, state_u=tick1["state_u"])) + want2 = np.asarray(ref.compute(x0, x_ref, eps2)).ravel() + np.testing.assert_allclose(tick2["out"], want2, rtol=1e-11, atol=1e-11) + + def test_graph_uses_plant_dynamics_and_only_u_recurs(self): + """Drift guard: the nonlinear op, the port layout, and the single state.""" + cg = _build_mppi_pendulum_graph() + ops = {node.op for node in cg.graph.nodes} + assert "sin" in ops, "the nonlinear rollout lost the plant's sin term" + for op in ("min", "matmul", "clip", "slice", "stack", "exp", "reshape"): + assert op in ops, f"nonlinear MPPI graph lost the {op!r} op" + + assert cg.outputs == ["out", "costs"] + # ``state_outputs`` is exactly what trace-time state detection found: + # the nominal plan, and nothing else (e.g. the tracking reference must + # not be promoted to a recurrent port). + assert cg.state_outputs == ["state_u"] + port_shapes = {n.attrs["name"]: n.shape for n in cg.graph.nodes if n.op == "input"} + assert port_shapes["epsilon"] == (MPPI_N, MPPI_K * MPPI_NL_DU) + assert port_shapes["state_u"] == (MPPI_K, MPPI_NL_DU) + assert port_shapes["current_state"] == (MPPI_NL_DX,) + assert port_shapes["target_state"] == (MPPI_NL_DX,) + + def test_node_count_is_independent_of_sample_count(self): + """The headline property: nodes track K*D_u, not the number of samples. + + A per-sample (looped) rollout would grow as ``N*K`` — the reason the + batched ``dynamics`` contract exists. Tracing the same policy with 4x + the samples must produce the identical graph. + """ + small = _build_mppi_pendulum_graph(N=6) + large = _build_mppi_pendulum_graph(N=24) + + def eps_shape(cg): + return next(n.shape for n in cg.graph.nodes if n.op == "input" and n.attrs["name"] == "epsilon") + + # The 4x batch really is in the graph — in the *shape* of the port, not + # in the number of nodes. + assert eps_shape(small) == (6, MPPI_K * MPPI_NL_DU) + assert eps_shape(large) == (24, MPPI_K * MPPI_NL_DU) + assert len(large.graph.nodes) == len(small.graph.nodes) + + def test_so_matches_interpreter_and_numpy(self, mppi_pendulum_so): + """.so, interpreter, and live numpy agree on the same seeded draws.""" + lib, cg = mppi_pendulum_so + n_out, n_state = output_split(cg) + assert n_state == MPPI_K * MPPI_NL_DU # the nominal plan recurs + assert n_out == MPPI_NL_DU + MPPI_N # out (D_u) + costs (N) + + ref = _mppi_pendulum_controller() + rng = np.random.default_rng(41) + max_u_err = 0.0 + for _ in range(10): + ref.reset() + x0, x_ref, eps = self._draw(rng) + feeds = self._feeds(x0, x_ref, eps) + out, state = step_so(lib, pack_arrays(cg, feeds), n_out, n_state) + traced = interpret(cg.graph, feeds) + want_u = np.asarray(ref.compute(x0, x_ref, eps)).ravel() + + # kernel vs its own interpreter (the tight tier), then vs numpy + np.testing.assert_allclose(out[:MPPI_NL_DU], traced["out"], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(state, np.asarray(traced["state_u"]).ravel(), rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(out[MPPI_NL_DU:], np.asarray(traced["costs"]).ravel(), rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(out[:MPPI_NL_DU], want_u, rtol=1e-11, atol=1e-11) + max_u_err = max(max_u_err, float(np.max(np.abs(out[:MPPI_NL_DU] - want_u)))) + assert max_u_err < 1e-11, f"nonlinear MPPI .so drifted from live numpy: {max_u_err:.3e}" + + def test_cabi_recurrence_matches_numpy(self, mppi_pendulum_so): + """The kernel's own state buffer reproduces a second live tick. + + The host feeds ``state_out`` straight back as the next tick's + ``state_u`` — no numpy state in the loop. + """ + lib, cg = mppi_pendulum_so + n_out, n_state = output_split(cg) + start, stop = state_slices(cg)["state_u"] + ref = _mppi_pendulum_controller() + rng = np.random.default_rng(43) + x0, x_ref, eps1 = self._draw(rng) + _, _, eps2 = self._draw(rng) + + out1, state1 = step_so(lib, pack_arrays(cg, self._feeds(x0, x_ref, eps1)), n_out, n_state) + want1 = np.asarray(ref.compute(x0, x_ref, eps1)).ravel() + np.testing.assert_allclose(out1[:MPPI_NL_DU], want1, rtol=1e-11, atol=1e-11) + + plan = state1[start:stop].reshape(MPPI_K, MPPI_NL_DU) + out2, state2 = step_so(lib, pack_arrays(cg, self._feeds(x0, x_ref, eps2, state_u=plan)), n_out, n_state) + want2 = np.asarray(ref.compute(x0, x_ref, eps2)).ravel() + np.testing.assert_allclose(out2[:MPPI_NL_DU], want2, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose(state2[start:stop], np.asarray(ref.u).ravel(), rtol=1e-11, atol=1e-11) class TestSolveQpOracle: @@ -1167,14 +1979,8 @@ def _run_closed_loop(lib, cg, case, ticks=100): est = case.estimator() ctrl = case.controller() - so_est = { - port: case.est_init.get(port, np.zeros(shape)).astype(np.float64).copy() - for port, _, shape in case.est_state_ports - } - live_est = { - attr: case.est_init.get(port, np.zeros(shape)).astype(np.float64).copy() - for port, attr, shape in case.est_state_ports - } + so_est = {port: case.est_init.get(port, np.zeros(shape)).astype(np.float64).copy() for port, _, shape in case.est_state_ports} + live_est = {attr: case.est_init.get(port, np.zeros(shape)).astype(np.float64).copy() for port, attr, shape in case.est_state_ports} so_ctrl = {port: np.zeros(shape) for port, _, shape in case.ctrl_state_ports} live_ctrl = {attr: np.zeros(shape) for _, attr, shape in case.ctrl_state_ports} u_prev_so = np.zeros(n_u) @@ -1253,9 +2059,7 @@ def test_so_matches_live_components(self, request, case): max_err, saw_saturation = _run_closed_loop(lib, cg, case) if case.sat_threshold is not None: assert saw_saturation, f"{case.name}: oracle never saturated — anti-windup path untested" - assert max_err < case.tol, ( - f"{case.name}: .so diverged from live components over 100 ticks: max abs err = {max_err:.3e}" - ) + assert max_err < case.tol, f"{case.name}: .so diverged from live components over 100 ticks: max abs err = {max_err:.3e}" def test_comptime_n_vars_mismatch_rejects_build(tmp_path): @@ -1430,9 +2234,7 @@ def test_drift_guard_nodes_match_emitted_table(self, manifests): for mn, zn in zip(manifest_nodes, zig_nodes): assert mn["vm_op"] == zn["vm_op"], f"node {mn['i']} op mismatch" assert mn["inputs"] == zn["inputs"], f"node {mn['i']} wiring mismatch" - assert mn["rows"] == zn["rows"] and mn["cols"] == zn["cols"], ( - f"node {mn['i']} shape mismatch" - ) + assert mn["rows"] == zn["rows"] and mn["cols"] == zn["cols"], f"node {mn['i']} shape mismatch" assert mn["aux"] == zn["aux"], f"node {mn['i']} aux mismatch" def test_lowering_is_deterministic(self, tmp_path):