From a6400e9d26b750245a7afdd1ba0f93c2cf980a10 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 12:30:23 -0400 Subject: [PATCH 01/21] feat: add abs, sign, pow, lt ops and config validation for SMC controller --- lab-notes/daily/2026-09-14.md | 5 + src/shinro/codegen/lower_zig.py | 1 + src/shinro/codegen/ops.py | 27 +++ src/shinro/codegen/trace_backend.py | 22 ++ src/shinro/codegen/tracing.py | 26 +++ src/shinro/controllers/smc.py | 166 +++++++++++-- src/shinro/runtime/graph_data.zig | 1 + src/shinro/runtime/graph_data_manifest.json | 8 - src/shinro/runtime/lower.zig | 24 +- src/shinro/utils/array_backend.py | 16 ++ tests/test_zig_lowering.py | 243 ++++++++++++++++++++ 11 files changed, 506 insertions(+), 33 deletions(-) create mode 100644 lab-notes/daily/2026-09-14.md diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md new file mode 100644 index 0000000..396a8e8 --- /dev/null +++ b/lab-notes/daily/2026-09-14.md @@ -0,0 +1,5 @@ +# Lab Notes — 2026-09-14 + +### 2026-09-14 16:30 UTC — initial + + diff --git a/src/shinro/codegen/lower_zig.py b/src/shinro/codegen/lower_zig.py index 9b0b644..10b6bc5 100644 --- a/src/shinro/codegen/lower_zig.py +++ b/src/shinro/codegen/lower_zig.py @@ -140,6 +140,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,") lines.append("};") lines.append("") lines.append("pub const Node = struct {") diff --git a/src/shinro/codegen/ops.py b/src/shinro/codegen/ops.py index 43a603a..17073dd 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,6 +227,22 @@ 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]])) diff --git a/src/shinro/codegen/trace_backend.py b/src/shinro/codegen/trace_backend.py index 90948a7..2eca5ac 100644 --- a/src/shinro/codegen/trace_backend.py +++ b/src/shinro/codegen/trace_backend.py @@ -174,6 +174,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 +205,12 @@ 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 argmax(self, x: Tracer) -> Tracer: # numpy argmax over the last axis collapses it to a scalar index. return self._emit("argmax", [x], ()) @@ -218,6 +233,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/controllers/smc.py b/src/shinro/controllers/smc.py index 392b9e6..755bea2 100644 --- a/src/shinro/controllers/smc.py +++ b/src/shinro/controllers/smc.py @@ -21,6 +21,7 @@ """ from dataclasses import dataclass +from typing import Any import numpy as np @@ -29,6 +30,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 +67,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" @@ -66,6 +91,17 @@ 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 :math:`|c^T g(x)|`, the + controllability denominator. **A deployment design parameter, not + a numerical constant** — set it above the smallest + :math:`|c^T g|` 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 +113,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() @@ -141,6 +179,13 @@ def compute(self, x, f_x, g_x): For scalar input (``c^T g`` is scalar), uses direct division. For vector input, solves the least-squares problem. + 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,). f_x: Drift dynamics :math:`f(x)` evaluated at x (n,). @@ -150,15 +195,26 @@ 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 for scalar input + (concrete/eager backends only; the traced path cannot raise). + NotImplementedError: If the multi-input least-squares branch is + traced — ``np.linalg.lstsq`` has no graph op. """ - 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 +222,81 @@ 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. + + Args: + num: ``s_dot_desired - c^T f`` (1,). + cg: ``c^T g`` (1,). - 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: + 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: least-squares solve of ``(c^T g) u = num``. + + ``np.linalg.lstsq`` has no graph op, so this branch is live-only — a + traced call raises instead of silently freezing a trace-time solve. - return u + Args: + num: ``s_dot_desired - c^T f`` (1,). + cg: ``c^T g`` (n_u,). + + Returns: + Control input (n_u,). + + Raises: + NotImplementedError: If called under tracing. + """ + concrete = self.bk.to_numpy(cg) + if not isinstance(concrete, np.ndarray): + raise NotImplementedError( + "SMC lowering supports a single control input (n_u == 1); the " + "multi-input least-squares branch uses np.linalg.lstsq, which " + "has no graph op. Lower n_u > 1 with the closed-form min-norm " + "u = cg^T (cg cg^T)^-1 (s_dot_desired - c^T f) — a follow-up." + ) + cg_np = np.asarray(concrete).reshape(1, -1) + rhs_np = np.asarray(self.bk.to_numpy(num)).reshape(-1)[:1] + u_np, _, _, _ = np.linalg.lstsq(cg_np, rhs_np, rcond=None) + return self.bk.from_numpy(u_np.flatten()) def reset(self): """No internal state to reset for SMC.""" @@ -199,6 +315,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 +335,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/runtime/graph_data.zig b/src/shinro/runtime/graph_data.zig index b8faea4..4ffeb4d 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, }; pub const Node = struct { diff --git a/src/shinro/runtime/graph_data_manifest.json b/src/shinro/runtime/graph_data_manifest.json index 58e05d5..f602e02 100644 --- a/src/shinro/runtime/graph_data_manifest.json +++ b/src/shinro/runtime/graph_data_manifest.json @@ -654,14 +654,6 @@ ] } ], - "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_outputs": [ { diff --git a/src/shinro/runtime/lower.zig b/src/shinro/runtime/lower.zig index 157afad..b4b5be5 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 @@ -253,7 +266,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,7 +297,7 @@ 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); @@ -303,6 +316,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..14174f2 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: ... diff --git a/tests/test_zig_lowering.py b/tests/test_zig_lowering.py index 48be2e8..feaecd4 100644 --- a/tests/test_zig_lowering.py +++ b/tests/test_zig_lowering.py @@ -33,6 +33,7 @@ from shinro.codegen.trace_node import trace_node from shinro.codegen.tracing import Graph 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 @@ -234,6 +235,89 @@ 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, +) -> 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. + + 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, 1)}) + 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) -> 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. + """ + x = rng.normal(0.0, 0.5, (2,)) + f_x = rng.normal(0.0, 0.5, (2,)) + while True: + g_x = rng.normal(0.0, 1.0, (2, 1)) + if abs(g_x[0, 0] + 2.0 * g_x[1, 0]) >= min_cg: + return {"x": x, "f_x": f_x, "g_x": g_x} + + @pytest.fixture(scope="session") def base_so(tmp_path_factory): """Build the .so from the base_tracking composed graph once per session.""" @@ -299,6 +383,37 @@ 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") + + +# 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 = { @@ -972,6 +1087,134 @@ def test_lowered_ops_match_interpreter(self, lowered_ops_so): ) +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_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 TestSolveQpOracle: """The .solve_qp VM op (codegen static solver) matches the interpreter. From ce09df6d5e5e9e87b1d7f81d67f6d4262d4969d7 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 13:00:20 -0400 Subject: [PATCH 02/21] feat: add SMC controller support and new ops (abs, sign, pow, lt, ne, stack, solve_qp) with Zig lowering updates and documentation enhancements --- docs/codegen.md | 47 ++++++++++++--- lab-notes/daily/2026-09-14.md | 84 ++++++++++++++++++++++++++ src/shinro/codegen/lower_zig.py | 4 +- src/shinro/codegen/trace_backend.py | 5 +- src/shinro/runtime/graph_data.zig | 87 +++++++++++++++++---------- src/shinro/runtime/lower.zig | 2 +- tests/test_op_shape_matrix.py | 61 +++++++++++++------ tests/test_zig_lowering.py | 92 ++++++++++++++--------------- 8 files changed, 270 insertions(+), 112 deletions(-) diff --git a/docs/codegen.md b/docs/codegen.md index 8949cb3..19b3d02 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`. | @@ -179,6 +179,35 @@ state_P (recurrent) ─────▶ Estimator (any state_* port the trace d `output` nodes are markers and are skipped; `compose` declares the combined outputs itself. +### 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-least-squares split keys off + `g_x`'s trace-time shape, not on traced values (`np.linalg.lstsq` has no + graph op, so the multi-input branch raises when traced). Dot products use + column-vector form because the tracer rejects 1D @ 1D. +- **Guards become data.** `c^T g` 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(abs(cg), eps)` then `u = where(cond, 0.0, u_raw)` — a fail-safe + zero command — and 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. @@ -209,9 +238,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`, `argmax`, `one_hot`, +`slice`, `solve_qp`. ## Lowering to Zig (shipped) @@ -347,11 +377,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`, `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/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md index 396a8e8..16d6b1d 100644 --- a/lab-notes/daily/2026-09-14.md +++ b/lab-notes/daily/2026-09-14.md @@ -2,4 +2,88 @@ ### 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 + diff --git a/src/shinro/codegen/lower_zig.py b/src/shinro/codegen/lower_zig.py index 10b6bc5..e88e994 100644 --- a/src/shinro/codegen/lower_zig.py +++ b/src/shinro/codegen/lower_zig.py @@ -250,9 +250,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( { diff --git a/src/shinro/codegen/trace_backend.py b/src/shinro/codegen/trace_backend.py index 2eca5ac..44d80f1 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,)) diff --git a/src/shinro/runtime/graph_data.zig b/src/shinro/runtime/graph_data.zig index 4ffeb4d..788b1af 100644 --- a/src/shinro/runtime/graph_data.zig +++ b/src/shinro/runtime/graph_data.zig @@ -2,11 +2,37 @@ // A ComposedGraph serialized as a comptime data table. pub const Op = enum { - cst, inp, out, matmul, add, sub, mul, div, ne, neg, - transpose, inv, reshape, clip, where_op, any, - copy, tanh, relu, exp, argmax, one_hot, slice, - sin, cos, stack, solve_qp, - abs, sign, pow, lt, + cst, + inp, + out, + matmul, + add, + sub, + mul, + div, + ne, + neg, + transpose, + inv, + reshape, + clip, + where_op, + any, + copy, + tanh, + relu, + exp, + argmax, + one_hot, + slice, + sin, + cos, + stack, + solve_qp, + abs, + sign, + pow, + lt, }; pub const Node = struct { @@ -35,43 +61,43 @@ pub const nodes = [_]Node{ .{ .op = .reshape, .inputs = &.{0}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .reshape, .inputs = &.{2}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{7, 3}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 7, 3 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 9, .vec = false }, - .{ .op = .matmul, .inputs = &.{9, 6}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .add, .inputs = &.{8, 10}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 9, 6 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{ 8, 10 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 18, .vec = false }, - .{ .op = .matmul, .inputs = &.{12, 4}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 12, 4 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 27, .vec = false }, - .{ .op = .matmul, .inputs = &.{13, 14}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 13, 14 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 36, .vec = false }, - .{ .op = .add, .inputs = &.{15, 16}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{ 15, 16 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 45, .vec = false }, - .{ .op = .matmul, .inputs = &.{18, 17}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 18, 17 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 54, .vec = false }, - .{ .op = .matmul, .inputs = &.{19, 20}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 19, 20 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 63, .vec = false }, - .{ .op = .add, .inputs = &.{21, 22}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{ 21, 22 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 72, .vec = false }, - .{ .op = .matmul, .inputs = &.{17, 24}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 17, 24 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .inv, .inputs = &.{23}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{25, 26}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 25, 26 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 81, .vec = false }, - .{ .op = .matmul, .inputs = &.{28, 11}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 28, 11 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 90, .vec = false }, - .{ .op = .matmul, .inputs = &.{30, 6}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .add, .inputs = &.{29, 31}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .sub, .inputs = &.{5, 32}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{27, 33}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .add, .inputs = &.{11, 34}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 30, 6 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{ 29, 31 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .sub, .inputs = &.{ 5, 32 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 27, 33 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{ 11, 34 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 99, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 108, .vec = false }, - .{ .op = .matmul, .inputs = &.{27, 37}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .sub, .inputs = &.{36, 38}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{39, 17}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 27, 37 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .sub, .inputs = &.{ 36, 38 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{ 39, 17 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .reshape, .inputs = &.{35}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, - .{ .op = .sub, .inputs = &.{1, 41}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, + .{ .op = .sub, .inputs = &.{ 1, 41 }, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 117, .vec = false }, - .{ .op = .matmul, .inputs = &.{43, 42}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, + .{ .op = .matmul, .inputs = &.{ 43, 42 }, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .clip, .inputs = &.{44}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .out, .inputs = &.{45}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .out, .inputs = &.{35}, .rows = 3, .cols = 1, .aux = 1, .vec = false }, @@ -83,9 +109,8 @@ pub const const_blob = [_]f64{ 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.47ae147ae147bp-6, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-6, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-6, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.47ae147ae147bp-7, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-7, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-7, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.999999999999ap-4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.999999999999ap-4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.999999999999ap-4, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.72a8f38fccafdp+4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.72a8f38fccafdp+4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.1e9b2675a6625p+4, }; -pub const clip_lo = [_]f64{-0x1.0000000000000p-1, -0x1.0000000000000p-1, -0x1.0000000000000p+0}; -pub const clip_hi = [_]f64{0x1.0000000000000p-1, 0x1.0000000000000p-1, 0x1.0000000000000p+0}; +pub const clip_lo = [_]f64{ -0x1.0000000000000p-1, -0x1.0000000000000p-1, -0x1.0000000000000p+0 }; +pub const clip_hi = [_]f64{ 0x1.0000000000000p-1, 0x1.0000000000000p-1, 0x1.0000000000000p+0 }; pub const output_offsets = [_]usize{0}; -pub const state_offsets = [_]usize{0, 3, 12}; - +pub const state_offsets = [_]usize{ 0, 3, 12 }; diff --git a/src/shinro/runtime/lower.zig b/src/shinro/runtime/lower.zig index b4b5be5..7ef766b 100644 --- a/src/shinro/runtime/lower.zig +++ b/src/shinro/runtime/lower.zig @@ -303,7 +303,7 @@ inline fn ew2(nodes: []const g.Node, node: g.Node, self_idx: usize, buf: *[g.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)]; diff --git a/tests/test_op_shape_matrix.py b/tests/test_op_shape_matrix.py index 1b9cb3a..e84e1d6 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)) @@ -182,23 +191,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], ()) @@ -208,8 +231,10 @@ def _pointwise_graph(g: Graph): outs["one_hot_4"] = g.emit("one_hot", [oh], (4,), depth=4) 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 +341,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_zig_lowering.py b/tests/test_zig_lowering.py index feaecd4..dc3aebd 100644 --- a/tests/test_zig_lowering.py +++ b/tests/test_zig_lowering.py @@ -136,8 +136,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=[], @@ -153,9 +163,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,)}, @@ -290,9 +298,7 @@ def _build_smc_graph( Returns: The traced, lowered-ready :class:`ComposedGraph`. """ - smc = _smc_controller( - smoother=smoother, phi=phi, alpha=alpha, controllability_eps=controllability_eps - ) + 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, 1)}) return ComposedGraph( graph=ng.graph, @@ -590,11 +596,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 @@ -714,11 +716,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 @@ -842,9 +840,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 @@ -865,53 +861,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 { @@ -920,6 +933,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) @@ -927,11 +941,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", @@ -1027,9 +1041,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) @@ -1082,9 +1094,7 @@ 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: @@ -1410,14 +1420,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) @@ -1496,9 +1500,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): @@ -1673,9 +1675,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): From 2289e937fed539233586f7a414d69178e7966543 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 13:30:17 -0400 Subject: [PATCH 03/21] style: reformat graph_data enums and array literals, tighten spacing --- lab-notes/daily/2026-09-14.md | 4 ++ src/shinro/runtime/graph_data.zig | 87 +++++++++++-------------------- 2 files changed, 35 insertions(+), 56 deletions(-) diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md index 16d6b1d..af2fcc4 100644 --- a/lab-notes/daily/2026-09-14.md +++ b/lab-notes/daily/2026-09-14.md @@ -87,3 +87,7 @@ 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 + + diff --git a/src/shinro/runtime/graph_data.zig b/src/shinro/runtime/graph_data.zig index 788b1af..4ffeb4d 100644 --- a/src/shinro/runtime/graph_data.zig +++ b/src/shinro/runtime/graph_data.zig @@ -2,37 +2,11 @@ // A ComposedGraph serialized as a comptime data table. pub const Op = enum { - cst, - inp, - out, - matmul, - add, - sub, - mul, - div, - ne, - neg, - transpose, - inv, - reshape, - clip, - where_op, - any, - copy, - tanh, - relu, - exp, - argmax, - one_hot, - slice, - sin, - cos, - stack, - solve_qp, - abs, - sign, - pow, - lt, + cst, inp, out, matmul, add, sub, mul, div, ne, neg, + transpose, inv, reshape, clip, where_op, any, + copy, tanh, relu, exp, argmax, one_hot, slice, + sin, cos, stack, solve_qp, + abs, sign, pow, lt, }; pub const Node = struct { @@ -61,43 +35,43 @@ pub const nodes = [_]Node{ .{ .op = .reshape, .inputs = &.{0}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .reshape, .inputs = &.{2}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 7, 3 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{7, 3}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 9, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 9, 6 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .add, .inputs = &.{ 8, 10 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{9, 6}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{8, 10}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 18, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 12, 4 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{12, 4}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 27, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 13, 14 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{13, 14}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 36, .vec = false }, - .{ .op = .add, .inputs = &.{ 15, 16 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{15, 16}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 45, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 18, 17 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{18, 17}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 54, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 19, 20 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{19, 20}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 63, .vec = false }, - .{ .op = .add, .inputs = &.{ 21, 22 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{21, 22}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 72, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 17, 24 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{17, 24}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .inv, .inputs = &.{23}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 25, 26 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{25, 26}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 81, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 28, 11 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{28, 11}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 90, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 30, 6 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .add, .inputs = &.{ 29, 31 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .sub, .inputs = &.{ 5, 32 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 27, 33 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, - .{ .op = .add, .inputs = &.{ 11, 34 }, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{30, 6}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{29, 31}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .sub, .inputs = &.{5, 32}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{27, 33}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, + .{ .op = .add, .inputs = &.{11, 34}, .rows = 3, .cols = 1, .aux = 0, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 99, .vec = false }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 108, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 27, 37 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .sub, .inputs = &.{ 36, 38 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 39, 17 }, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{27, 37}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .sub, .inputs = &.{36, 38}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, + .{ .op = .matmul, .inputs = &.{39, 17}, .rows = 3, .cols = 3, .aux = 0, .vec = false }, .{ .op = .reshape, .inputs = &.{35}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, - .{ .op = .sub, .inputs = &.{ 1, 41 }, .rows = 3, .cols = 1, .aux = 0, .vec = true }, + .{ .op = .sub, .inputs = &.{1, 41}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .cst, .inputs = &.{}, .rows = 3, .cols = 3, .aux = 117, .vec = false }, - .{ .op = .matmul, .inputs = &.{ 43, 42 }, .rows = 3, .cols = 1, .aux = 0, .vec = true }, + .{ .op = .matmul, .inputs = &.{43, 42}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .clip, .inputs = &.{44}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .out, .inputs = &.{45}, .rows = 3, .cols = 1, .aux = 0, .vec = true }, .{ .op = .out, .inputs = &.{35}, .rows = 3, .cols = 1, .aux = 1, .vec = false }, @@ -109,8 +83,9 @@ pub const const_blob = [_]f64{ 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.47ae147ae147bp-6, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-6, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-6, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.47ae147ae147bp-7, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-7, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.47ae147ae147bp-7, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.999999999999ap-4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.999999999999ap-4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.999999999999ap-4, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.0000000000000p+0, 0x1.72a8f38fccafdp+4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.72a8f38fccafdp+4, 0x0.0p+0, 0x0.0p+0, 0x0.0p+0, 0x1.1e9b2675a6625p+4, }; -pub const clip_lo = [_]f64{ -0x1.0000000000000p-1, -0x1.0000000000000p-1, -0x1.0000000000000p+0 }; -pub const clip_hi = [_]f64{ 0x1.0000000000000p-1, 0x1.0000000000000p-1, 0x1.0000000000000p+0 }; +pub const clip_lo = [_]f64{-0x1.0000000000000p-1, -0x1.0000000000000p-1, -0x1.0000000000000p+0}; +pub const clip_hi = [_]f64{0x1.0000000000000p-1, 0x1.0000000000000p-1, 0x1.0000000000000p+0}; pub const output_offsets = [_]usize{0}; -pub const state_offsets = [_]usize{ 0, 3, 12 }; +pub const state_offsets = [_]usize{0, 3, 12}; + From 2aae56b0cf5fdcfffe9941a127f675205255006b Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 19:30:19 -0400 Subject: [PATCH 04/21] feat: implement min-norm multi-input control in SlidingModeController and update documentation --- docs/codegen.md | 25 +++++---- lab-notes/daily/2026-09-14.md | 4 ++ src/shinro/controllers/smc.py | 95 +++++++++++++++++++++++------------ tests/test_controllers.py | 65 ++++++++++++++++++++++++ 4 files changed, 147 insertions(+), 42 deletions(-) diff --git a/docs/codegen.md b/docs/codegen.md index 19b3d02..cc2a9a8 100644 --- a/docs/codegen.md +++ b/docs/codegen.md @@ -192,15 +192,22 @@ relying on `build_composed_graph`'s role map. Two behavioral conventions make it lowerable at all: -- **Shape-driven branches.** The scalar-vs-least-squares split keys off - `g_x`'s trace-time shape, not on traced values (`np.linalg.lstsq` has no - graph op, so the multi-input branch raises when traced). Dot products use - column-vector form because the tracer rejects 1D @ 1D. -- **Guards become data.** `c^T g` 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(abs(cg), eps)` then `u = where(cond, 0.0, u_raw)` — a fail-safe - zero command — and publishes `healthy = 1 - cond` through +- **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 diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md index af2fcc4..21f046e 100644 --- a/lab-notes/daily/2026-09-14.md +++ b/lab-notes/daily/2026-09-14.md @@ -91,3 +91,7 @@ flag at all — SMC's compiled form is now the more diagnosable of the two. ### 2026-09-14 17:30 UTC — update + +### 2026-09-14 23:30 UTC — update + + diff --git a/src/shinro/controllers/smc.py b/src/shinro/controllers/smc.py index 755bea2..1052985 100644 --- a/src/shinro/controllers/smc.py +++ b/src/shinro/controllers/smc.py @@ -91,17 +91,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 :math:`|c^T g(x)|`, the - controllability denominator. **A deployment design parameter, not - a numerical constant** — set it above the smallest - :math:`|c^T g|` 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. + 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. """ @@ -177,7 +179,11 @@ 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), @@ -195,10 +201,8 @@ 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 - (concrete/eager backends only; the traced path cannot raise). - NotImplementedError: If the multi-input least-squares branch is - traced — ``np.linalg.lstsq`` has no graph op. + 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.ravel(self.bk.array(x)) f_x = self.bk.ravel(self.bk.array(f_x)) @@ -270,10 +274,21 @@ def _scalar_control(self, num, cg): return self.bk.where(fault, self.bk.zeros_like(u_raw), u_raw) def _vector_control(self, num, cg): - """Multi-input law: least-squares solve of ``(c^T g) u = num``. - - ``np.linalg.lstsq`` has no graph op, so this branch is live-only — a - traced call raises instead of silently freezing a trace-time solve. + """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,). @@ -283,20 +298,34 @@ def _vector_control(self, num, cg): Control input (n_u,). Raises: - NotImplementedError: If called under tracing. + RuntimeError: If ``‖c^T g‖ < controllability_eps`` and ``cg`` + is concrete. """ - concrete = self.bk.to_numpy(cg) - if not isinstance(concrete, np.ndarray): - raise NotImplementedError( - "SMC lowering supports a single control input (n_u == 1); the " - "multi-input least-squares branch uses np.linalg.lstsq, which " - "has no graph op. Lower n_u > 1 with the closed-form min-norm " - "u = cg^T (cg cg^T)^-1 (s_dot_desired - c^T f) — a follow-up." - ) - cg_np = np.asarray(concrete).reshape(1, -1) - rhs_np = np.asarray(self.bk.to_numpy(num)).reshape(-1)[:1] - u_np, _, _, _ = np.linalg.lstsq(cg_np, rhs_np, rcond=None) - return self.bk.from_numpy(u_np.flatten()) + 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)) + + # 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.""" diff --git a/tests/test_controllers.py b/tests/test_controllers.py index f02bf47..b3a8481 100644 --- a/tests/test_controllers.py +++ b/tests/test_controllers.py @@ -443,6 +443,71 @@ 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 From 589819acace83d95dc663c30e5edcadef0baf774 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 20:00:19 -0400 Subject: [PATCH 05/21] feat: add multi-input SMC support with closed-form vector control and comprehensive tests --- lab-notes/daily/2026-09-14.md | 64 +++++++++++++++-- tests/test_controllers.py | 129 ++++++++++++++++++++++++++++------ tests/test_zig_lowering.py | 88 +++++++++++++++++++++-- 3 files changed, 250 insertions(+), 31 deletions(-) diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md index 21f046e..4d19d41 100644 --- a/lab-notes/daily/2026-09-14.md +++ b/lab-notes/daily/2026-09-14.md @@ -86,12 +86,68 @@ 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 + diff --git a/tests/test_controllers.py b/tests/test_controllers.py index b3a8481..4830e89 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]) @@ -452,6 +482,7 @@ def test_smc_multi_input_matches_min_norm_reference(self, bk): 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): @@ -478,6 +509,7 @@ def test_smc_multi_input_matches_min_norm_reference(self, bk): 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]) @@ -501,6 +533,7 @@ def test_smc_multi_input_min_norm_is_shortest(self, bk): 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]) @@ -511,6 +544,7 @@ def test_smc_multi_input_loss_of_controllability(self, bk): 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]) @@ -524,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]) @@ -540,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]) @@ -557,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]) @@ -574,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]) @@ -591,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]) @@ -608,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]) @@ -623,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]) @@ -638,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]) @@ -648,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]) @@ -658,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]) @@ -668,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]) @@ -678,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() @@ -685,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 @@ -701,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 @@ -716,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), @@ -733,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"): @@ -751,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])) @@ -902,6 +957,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 @@ -912,6 +968,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) @@ -937,10 +994,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])) @@ -950,10 +1013,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])) @@ -963,10 +1032,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) @@ -975,12 +1050,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]) @@ -998,10 +1079,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_zig_lowering.py b/tests/test_zig_lowering.py index dc3aebd..10e34a8 100644 --- a/tests/test_zig_lowering.py +++ b/tests/test_zig_lowering.py @@ -272,6 +272,7 @@ def _build_smc_graph( 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. @@ -292,14 +293,17 @@ def _build_smc_graph( 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 + 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, 1)}) + 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"], @@ -309,18 +313,20 @@ def _build_smc_graph( ) -def _smc_rand_inputs(rng: np.random.Generator, min_cg: float = 0.2) -> dict[str, np.ndarray]: - """Random SMC inputs with ``|c^T g|`` kept above the guard. +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, 1)) - if abs(g_x[0, 0] + 2.0 * g_x[1, 0]) >= min_cg: + 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} @@ -400,6 +406,17 @@ def smc_so(tmp_path_factory): 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") + + # 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. @@ -1165,6 +1182,65 @@ def test_lost_controllability_is_failsafe_and_flagged(self, smc_so, g_x): 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 From b7bdb945fbc72634fe4aca0c0b1199de8b4c1112 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 20:30:19 -0400 Subject: [PATCH 06/21] docs: defer multi-surface SMC decision and clarify single-surface scope in documentation --- lab-notes/daily/2026-09-14.md | 56 +++++++++++++++++++++++++++++++++++ src/shinro/controllers/smc.py | 21 +++++++++++++ 2 files changed, 77 insertions(+) diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md index 4d19d41..a549354 100644 --- a/lab-notes/daily/2026-09-14.md +++ b/lab-notes/daily/2026-09-14.md @@ -150,4 +150,60 @@ 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 + diff --git a/src/shinro/controllers/smc.py b/src/shinro/controllers/smc.py index 1052985..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" @@ -80,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. From bf792d1b672c159c052994a14f1ab11bd72052d6 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Mon, 14 Sep 2026 21:00:21 -0400 Subject: [PATCH 07/21] style: trim extra blank lines in daily note markdown file --- lab-notes/daily/2026-09-14.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lab-notes/daily/2026-09-14.md b/lab-notes/daily/2026-09-14.md index a549354..c63f54c 100644 --- a/lab-notes/daily/2026-09-14.md +++ b/lab-notes/daily/2026-09-14.md @@ -203,7 +203,8 @@ is an extension, not a restart. ("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 + From d012e7a4e7deac8c2c3e6550864bd7a3d658469b Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 12:18:08 -0400 Subject: [PATCH 08/21] docs: add auto-generated API reference site Adds an mkdocs-material site whose API reference is generated from source rather than hand-maintained, plus a Pages deploy workflow. The reference is derived from each subpackage's __all__: scripts/gen_api.py walks __all__ -> one page per subpackage + the nav. Subpackages without __all__ (components.py, utils/) fall back to an AST scan of their source files. Adding an export is therefore all it takes for it to appear in the docs; no config, nav, or symbol-list edit is needed. Verified by injecting a class into controllers/ and observing the page grow from 7 to 8 exports. scripts/sphinx_compat.py is a griffe extension handling the Sphinx-flavored docstring markup used throughout the source (:class:, :meth:, :func:, :math:, and .. math:: blocks), which griffe's Google parser does not understand and would otherwise render as literal text. Without it the docs build still succeeds but every formula is broken, so the key must stay under `options:` in mkdocs.yml rather than at handler level. Generated output (docs/reference/, docs/SUMMARY.md, site/) is gitignored; CI regenerates it on every run. Notes: - The workflow deliberately omits --strict: griffe reports the unannotated public parameters as warnings (204 today), so --strict fails out of the box. Worth adding once annotation coverage improves. - 53% of public symbols have docstrings, so roughly 281 render as bare signatures. This is a source-coverage limit, not a tooling one. - Prose pages remain hand-written and are ordered in docs/_nav_prose.md. Commands: make docs, make docs-serve, make docs-build. --- lab-notes/daily/2026-09-15.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 lab-notes/daily/2026-09-15.md diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md new file mode 100644 index 0000000..8bfd0df --- /dev/null +++ b/lab-notes/daily/2026-09-15.md @@ -0,0 +1,5 @@ +# Lab Notes — 2026-09-15 + +### 2026-09-15 16:18 UTC — initial + + From 3f3419609886f36882a9aa6978aa2b5c27d5745e Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 14:30:36 -0400 Subject: [PATCH 09/21] feat: add min reduction op with Zig lowering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MPPI's softmax shift (beta = min(costs)) is the one 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. Add it end-to-end so MPPI can lower. - codegen/ops.py: @register_op("min") — np.min with axis as a node attr (None -> 0-d scalar, 0/1 -> that axis of a 2-D input). The interpreter half, i.e. the oracle the .so is checked against. - codegen/trace_backend.py: TraceBackend.min(x, axis=None), output shape fixed at trace time. - runtime/linalg.zig: min_all / min_axis0 / min_axis1 (the argmax precedent: kernel math lives here, the VM only dispatches). Each is NaN-aware to mirror numpy's NaN-wins semantics — a plain `v < best` silently ignores NaNs. - runtime/lower.zig: .min dispatch arm; axis rides in node.aux (0 = None, 1 = axis 0, 2 = axis 1), like slice's start offset. - codegen/lower_zig.py: min joins the generated Op enum; _node_vm_info maps axis -> aux. - tests/test_op_shape_matrix.py: 5 cells in the pointwise group — 1-D full, 2-D full, 2-D axis 0, 2-D axis 1, and a div-by-zero NaN-propagation cell. The rank-3 lowering guard that forces MPPI's (N,K,D_u) epsilon to flatten at the port boundary already shipped in afe2e2c; verified, not re-implemented. Verified: make test 1135 passed/7 skipped; make test-zig 64 passed/2 skipped; op/shape matrix 5 passed; make lint 0 errors. --- lab-notes/daily/2026-09-15.md | 47 +++++++++++++++++++++++++++++ src/shinro/codegen/lower_zig.py | 8 ++++- src/shinro/codegen/ops.py | 14 +++++++++ src/shinro/codegen/trace_backend.py | 16 ++++++++++ src/shinro/runtime/graph_data.zig | 2 +- src/shinro/runtime/linalg.zig | 43 ++++++++++++++++++++++++++ src/shinro/runtime/lower.zig | 18 +++++++++++ tests/test_op_shape_matrix.py | 10 ++++++ 8 files changed, 156 insertions(+), 2 deletions(-) diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index 8bfd0df..926f3da 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -2,4 +2,51 @@ ### 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), and the manifest's `provenance` block is restored +(it was dropped by `e3c38a1`'s `make zig-gen`; `0bf27b3` had introduced it). + +**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:30 UTC — update + diff --git a/src/shinro/codegen/lower_zig.py b/src/shinro/codegen/lower_zig.py index e88e994..1d54349 100644 --- a/src/shinro/codegen/lower_zig.py +++ b/src/shinro/codegen/lower_zig.py @@ -140,7 +140,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,") + lines.append(" abs, sign, pow, lt, min,") lines.append("};") lines.append("") lines.append("pub const Node = struct {") @@ -386,6 +386,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/ops.py b/src/shinro/codegen/ops.py index 17073dd..a3ba8c4 100644 --- a/src/shinro/codegen/ops.py +++ b/src/shinro/codegen/ops.py @@ -248,6 +248,20 @@ def _argmax(node: Node, values: dict[int, np.ndarray], inputs: dict[str, np.ndar 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/trace_backend.py b/src/shinro/codegen/trace_backend.py index 44d80f1..02b1973 100644 --- a/src/shinro/codegen/trace_backend.py +++ b/src/shinro/codegen/trace_backend.py @@ -208,6 +208,22 @@ def abs(self, x: Tracer) -> Tracer: 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], ()) diff --git a/src/shinro/runtime/graph_data.zig b/src/shinro/runtime/graph_data.zig index 4ffeb4d..42681a8 100644 --- a/src/shinro/runtime/graph_data.zig +++ b/src/shinro/runtime/graph_data.zig @@ -6,7 +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, + abs, sign, pow, lt, min, }; pub const Node = struct { 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 7ef766b..3b123f5 100644 --- a/src/shinro/runtime/lower.zig +++ b/src/shinro/runtime/lower.zig @@ -205,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]); diff --git a/tests/test_op_shape_matrix.py b/tests/test_op_shape_matrix.py index e84e1d6..b93e1f5 100644 --- a/tests/test_op_shape_matrix.py +++ b/tests/test_op_shape_matrix.py @@ -230,6 +230,16 @@ 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"), From 24a007a782b8a0aba375aceb08c8c4972cdbafe9 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 14:44:57 -0400 Subject: [PATCH 10/21] chore: added pi specific settings --- .pi-lens.json | 3 +++ lab-notes/daily/2026-09-15.md | 4 ++++ 2 files changed, 7 insertions(+) create mode 100644 .pi-lens.json 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/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index 926f3da..5c8834b 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -50,3 +50,7 @@ None/0/1 raises. ### 2026-09-15 18:30 UTC — update + +### 2026-09-15 18:44 UTC — update + + From 3083225cc7f902cda3e3b85eea406430cde42c9f Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 14:52:39 -0400 Subject: [PATCH 11/21] refactor: make the batched quadratic cost traceable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BatchedDynamicsAdapter._quad_form used bk.sum(..., axis=1), but the adapter holds the plant's backend rather than the controller's TraceBackend, so the sum never sees a traced operand. Rewrite both branches as contractions built from operators (which the Tracer overloads intercept, lifting the concrete operands as consts): - diagonal W: (z*z) @ W — W itself is the contraction vector, matvec (N,D) @ (D,) -> (N,) - full W: (z * (z @ W.T)) @ ones — the row-wise dot product contracted with a ones vector A sum over an axis is a matmul identity, so no new VM op is needed. The LTI dynamics_fn was already operator-only; the nonlinear _integrate path is untouched and stays eager-only (the phase-6 graph build refuses it). Note: the Tracer's _broadcast_shape rejects numpy's rank-differing elementwise broadcast ((N,D) * (D,)), even though the VM's bcast_flat supports it — so a rank-differing operand must be reshaped same-rank or absorbed into a contraction. The diagonal branch does the latter. Verified: adapter tests 15 passed/3 skipped; eager MPPI 36 passed/2 skipped; both branches traced through interpret() vs numpy (diag exact, full 4.4e-16); make test 1135 passed/7 skipped; make lint 0 errors. --- lab-notes/daily/2026-09-15.md | 52 +++++++++++++++++++++++++++-- src/shinro/utils/batched_adapter.py | 18 ++++++++-- 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index 5c8834b..5525b8c 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -3,7 +3,6 @@ ### 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. @@ -40,13 +39,56 @@ boundary — verified, not re-implemented. Reconfirmed the guard test passes. `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), and the manifest's `provenance` block is restored -(it was dropped by `e3c38a1`'s `make zig-gen`; `0bf27b3` had introduced it). +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 @@ -54,3 +96,7 @@ None/0/1 raises. ### 2026-09-15 18:44 UTC — update + +### 2026-09-15 18:52 UTC — update + + diff --git a/src/shinro/utils/batched_adapter.py b/src/shinro/utils/batched_adapter.py index fd2bd78..c7492aa 100644 --- a/src/shinro/utils/batched_adapter.py +++ b/src/shinro/utils/batched_adapter.py @@ -147,6 +147,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 +164,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 From e53729fbfefcb1cc79f4f75005aa972a43b62087 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 15:51:48 -0400 Subject: [PATCH 12/21] feat: make MPPI's compute trace-safe for lowering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewrite MPPIController.compute so one implementation serves eager numpy, eager torch, and the tracing backend — the prerequisite for lowering the LTI MPPI rollout to Zig. - Sampling stays host-side: new optional epsilon arg, shape (N, K*D_u) sample-major. None draws from self._rng (unchanged eager path); the traced call receives it as a graph input port, so the branch is a trace-time constant and the RNG is never touched while tracing. - Everything after sampling goes through self.bk: the initial state batches via a same-rank broadcast, every sum is a matmul identity (control penalty = ones contraction, weighted update = one (1,N) @ (N, K*D_u), softmax normalizer = (1,N) @ ones), and the only genuine reduction is beta = min(costs). - self.u is rebound, not mutated in place (detect_state keys off reassignment); the receding-horizon shift is K row-slices + a stack. - The nominal sequence and host-supplied epsilon are bridged to the backend up front: numpy + torch raises (only torch + numpy works). - costs is published via emit_named_output as a graph output port (SMC's healthy precedent); the tracking reference lives in a 1-element list and is passed as a (1, D_x) row, so neither is mis-detected as recurrent state. The tracer detects exactly one state attr: u. - ArrayBackend.min (numpy/torch) and a TorchBackend.clip fix for numpy bounds, both needed by the eager path now that clipping happens in backend space. - _as_float (mirroring smc.py) replaces bare float(dt)/float(temperature). Verified: MPPI 38 passed/2 skipped (both backends), adapter 15/3, trace smoke bit-exact vs eager numpy (u/state_u err 0.0, outputs costs/out/state_u), make test 1137 passed/7 skipped, make lint 0 errors. --- lab-notes/daily/2026-09-15.md | 106 +++++++++++++- src/shinro/controllers/mppi.py | 222 +++++++++++++++++++++--------- src/shinro/utils/array_backend.py | 20 +++ tests/test_controllers.py | 24 ++++ 4 files changed, 303 insertions(+), 69 deletions(-) diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index 5525b8c..9c82225 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -3,6 +3,7 @@ ### 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. @@ -91,12 +92,109 @@ 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 + diff --git a/src/shinro/controllers/mppi.py b/src/shinro/controllers/mppi.py index 041fa27..5e85e63 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,22 @@ 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 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 @@ -200,20 +234,21 @@ def attach_plant(self, plant, Q: Any | None = None, R: Any | None = None): 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) + 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 +256,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 +279,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() + x_current = dynamics_fn(x_current, v_k, self.dt) - 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) - - 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/utils/array_backend.py b/src/shinro/utils/array_backend.py index 14174f2..d319436 100644 --- a/src/shinro/utils/array_backend.py +++ b/src/shinro/utils/array_backend.py @@ -173,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: ... @@ -378,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]) @@ -548,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): @@ -617,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/tests/test_controllers.py b/tests/test_controllers.py index 4830e89..b6f68ec 100644 --- a/tests/test_controllers.py +++ b/tests/test_controllers.py @@ -893,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 From ed5c40a9719a3b1cfa2bc125f9e15cfbfbd4dd1a Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 15:51:48 -0400 Subject: [PATCH 13/21] docs: add an MPPI demo for model wiring and lowering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit demos/demo_mppi.py walks the four ways to give MPPI its dynamics_fn/cost_fn (constructor injection, attribute injection, attach_plant, from_config + injection), the host-supplied epsilon contract a lowered kernel relies on, and an end-to-end trace of a hand-written model through trace_node + interpret() vs live numpy (u max err 0.0). It closes with the trace-safety rules for a lowerable model. MuJoCo-free — runs on the default install. Also records the first comptime-scale data point: a single-axis MPPI graph (N=200, K=15) traces to 388 nodes. --- demos/demo_mppi.py | 243 ++++++++++++++++++++++++++++++++++ lab-notes/daily/2026-09-15.md | 4 + 2 files changed, 247 insertions(+) create mode 100644 demos/demo_mppi.py diff --git a/demos/demo_mppi.py b/demos/demo_mppi.py new file mode 100644 index 0000000..049138c --- /dev/null +++ b/demos/demo_mppi.py @@ -0,0 +1,243 @@ +# 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", + "dynamics must be LTI for lowering; nonlinear rollouts stay eager-only", + ), + start=1, + ): + print(f" {i}. {rule}") + print("\n attach_plant's LTI path already follows all of these.") + + +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/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index 9c82225..b894ee5 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -198,3 +198,7 @@ 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:51 UTC — update + + From 89ccbd7918709505b023bc7d861cc8acba9ca770 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 16:54:26 -0400 Subject: [PATCH 14/21] feat: lower MPPI to Zig end-to-end, and record kernel sizes as a metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MPPI is the first lowered controller whose Gaussian sampling stays on the host: the perturbations arrive through a free `epsilon` port, which is what makes three-way parity checkable at all — same draw in, identical u/state/costs out. Lowering (phases 6-7): - lower_zig: clip bounds now broadcast via np.broadcast_to. MPPI clips an (N, D_u) sample batch with per-channel (D_u,) limits, which the old same-size-or-scalar clip-blob expansion refused. `vals.tolist()` replaces float(v) (same result, and avoids the linter's bare-float() rule). - _build_mppi_graph + TestMppiOracle: a standalone graph (epsilon + recurrent state_u in; out + costs out), interpreter-only parity / costs-port / two-tick-recurrence / structure tests, then a mppi_so fixture adding .so-vs-interpreter-vs-numpy parity and a C-ABI recurrence test. - Measured: node count tracks K*D_u (the unrolled rollout), not N; the VM stack buffer tracks N*K (~50 f64 per sample-step). At the shipped config (N=200, K=15) that extrapolates to ~1.1 MiB and ~3 min of ReleaseFast compile; the production-scale check (N=100, K=15) is float-exact (~1e-16 vs numpy). Kernel size metric: - graph_data_manifest.json now self-reports byte sizes: buf_bytes, const_blob_bytes, clip_blob_len/bytes, input/output/state_bytes, and a `bytes` entry per port (additive; no key removed). - `make measure-kernels` (shinro/codegen/measure.py + a scripts/ shim) reports the C-ABI host buffers, the VM stack buffer, and, with BUILD=1, the artifact bytes and compile cost across a D_x/D_u/N/K sweep. measure.graph_metrics reads the manifest fields rather than recomputing them. - Deliberately not a CI gate: the document includes a wall-clock compile time, so it is a measurement sample, not a diffable audit record. Verified: make test 1154 passed / 7 skipped; make test-zig 70 passed / 2 skipped; tests/test_measure.py 11 passed; make lint 0 errors. --- AGENTS.md | 9 +- Makefile | 14 +- docs/testing.md | 5 +- lab-notes/daily/2026-09-15.md | 207 +++++++++++++++- scripts/measure_kernels.py | 12 + src/shinro/codegen/lower_zig.py | 54 +++-- src/shinro/codegen/measure.py | 252 ++++++++++++++++++++ src/shinro/runtime/graph_data_manifest.json | 16 ++ tests/test_measure.py | 118 +++++++++ tests/test_op_shape_matrix.py | 6 + tests/test_zig_lowering.py | 251 +++++++++++++++++++ 11 files changed, 922 insertions(+), 22 deletions(-) create mode 100644 scripts/measure_kernels.py create mode 100644 src/shinro/codegen/measure.py create mode 100644 tests/test_measure.py 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/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-15.md b/lab-notes/daily/2026-09-15.md index b894ee5..f0562e0 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -197,8 +197,209 @@ 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: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 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/lower_zig.py b/src/shinro/codegen/lower_zig.py index 1d54349..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 @@ -239,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 @@ -265,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, 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/runtime/graph_data_manifest.json b/src/shinro/runtime/graph_data_manifest.json index f602e02..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,8 +656,10 @@ "reshape", "sub" ], + "output_bytes": 24, "outputs": [ { + "bytes": 24, "name": "u", "shape": [ 3 @@ -655,8 +667,10 @@ } ], "solve_qp": null, + "state_bytes": 120, "state_outputs": [ { + "bytes": 24, "name": "state_x_hat", "shape": [ 3, @@ -664,6 +678,7 @@ ] }, { + "bytes": 72, "name": "state_P", "shape": [ 3, @@ -671,6 +686,7 @@ ] }, { + "bytes": 24, "name": "state_u_prev", "shape": [ 3 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 b93e1f5..c8d64cb 100644 --- a/tests/test_op_shape_matrix.py +++ b/tests/test_op_shape_matrix.py @@ -153,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 diff --git a/tests/test_zig_lowering.py b/tests/test_zig_lowering.py index 10e34a8..3cb58ba 100644 --- a/tests/test_zig_lowering.py +++ b/tests/test_zig_lowering.py @@ -32,6 +32,7 @@ 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 @@ -330,6 +331,82 @@ def _smc_rand_inputs(rng: np.random.Generator, min_cg: float = 0.2, n_u: int = 1 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"], + ) + + @pytest.fixture(scope="session") def base_so(tmp_path_factory): """Build the .so from the base_tracking composed graph once per session.""" @@ -417,6 +494,20 @@ def smc_multi_so(tmp_path_factory): 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") + + # 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. @@ -1301,6 +1392,166 @@ def test_controllability_eps_is_baked_from_config(self, tmp_path): 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 TestSolveQpOracle: """The .solve_qp VM op (codegen static solver) matches the interpreter. From c834aad71d7d7eeb7e882f6ff5016fed782c5de0 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 17:14:30 -0400 Subject: [PATCH 15/21] feat: compile MPPI scenarios end to end (host-supplied epsilon) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `make compile` produced deployable kernels for LQR/PID/MPC but not MPPI: compose() refuses unmapped controller inputs, so the sampled-perturbation port had nowhere to go. The fix is 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=(...)) 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 is byte-for-byte unchanged; undeclared names still raise. - build_composed_graph(..., plant=None) attaches the plant when the controller supports it (the same attach_plant 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. - tests/integration/scenarios/mppi_compile.toml: KF + MPPI on holonomic_base, compile-ready (distinct from the sim-backed mppi_base_tracking.toml). End to end: 493 nodes; inputs [y, x_ref, u_prev, state_x_hat, state_P, state_u, epsilon]; oracle B (.so vs interpret) max abs err 6.4e-15; deployment record stamped and verified. The host now owns a 72,000-byte epsilon port. Tests (+5): TestComposeHostInputs (declared shape, appended last, not recurrent, an undeclared name still raises, and a real graph input fed to the interpreter); a build_composed_graph MPPI+plant test; and a hypothetical third estimator (an ad-hoc registered EMA estimator) composing with MPPI and matching the live loop tick-for-tick — proving the pipeline is estimator-agnostic. Verified: make test 1159 passed / 7 skipped; make test-zig 70 passed / 2 skipped; make lint 0 errors. --- lab-notes/daily/2026-09-15.md | 53 ++++++ src/shinro/codegen/build.py | 14 ++ src/shinro/codegen/compose.py | 27 ++- src/shinro/codegen/scenario_gen.py | 2 + src/shinro/controllers/mppi.py | 16 ++ tests/integration/scenarios/mppi_compile.toml | 32 ++++ tests/test_codegen_build.py | 162 +++++++++++++++++- tests/test_codegen_compose.py | 66 +++++++ 8 files changed, 368 insertions(+), 4 deletions(-) create mode 100644 tests/integration/scenarios/mppi_compile.toml diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index f0562e0..c2c703a 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -402,4 +402,57 @@ passed / 7 skipped (+2); ``make test-zig`` 70 passed / 2 skipped; ``make lint`` ### 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 + 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..e782c11 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. @@ -106,6 +107,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. @@ -196,6 +202,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 +212,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 +229,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,6 +266,12 @@ 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) # The controller's output (u) — clip if input_limits provided, then emit. @@ -337,7 +358,7 @@ def compose( return ComposedGraph( graph=combined, - inputs=["y", "x_ref", "u_prev"] + state_ports + ctrl_state_ports, + inputs=["y", "x_ref", "u_prev"] + state_ports + ctrl_state_ports + free_input_names, outputs=["u"], state_inputs=emitted_state_ports + emitted_ctrl_state_ports + ["u_prev"], state_outputs=emitted_state_ports + emitted_ctrl_state_ports + ["state_u_prev"], 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/controllers/mppi.py b/src/shinro/controllers/mppi.py index 5e85e63..59aabda 100644 --- a/src/shinro/controllers/mppi.py +++ b/src/shinro/controllers/mppi.py @@ -198,6 +198,22 @@ def __init__( # 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. 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/test_codegen_build.py b/tests/test_codegen_build.py index 97b9b8d..ec8660f 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.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,160 @@ 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) + assert cg.outputs == ["u"] + + +# ─── 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..e250549 100644 --- a/tests/test_codegen_compose.py +++ b/tests/test_codegen_compose.py @@ -614,6 +614,72 @@ 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,) + + # ─── Test 10: compose error paths ────────────────────────────────────────── From 8fc9f6c558207ed9ce1b3e88b04a220c3cee81b0 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Tue, 15 Sep 2026 17:28:27 -0400 Subject: [PATCH 16/21] fix(compose): forward subgraph diagnostics so composed kernels stay observable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A composed binary 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 deployed kernel therefore had no health signal and no per-tick cost to watch for divergence. `compose._forward_diagnostics` (called after each subgraph merge) now 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 (`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 random inputs, max abs err 1.137e-13; deployment record stamped and verified. LQR's layout is unchanged (['u']). Tests (+5): TestComposeDiagnostics — forwarded after u in estimator→controller order, a real interpreter-visible port, no-diagnostics leaves the layout alone, duplicate names raise, and a diagnostic named u raises. The MPPI build test also asserts cg.outputs == ['u', 'costs'] and `interpret` returning costs (200,). Docs: docs/codegen.md claimed subgraph output nodes "are skipped" — corrected to document the out/state_* explicit wiring plus diagnostic forwarding, and the `host_inputs` free-port mechanism from the previous batch. Verified: make test 1164 passed / 7 skipped; make test-zig 70 passed / 2 skipped; make lint 0 errors. --- docs/codegen.md | 20 ++++++--- lab-notes/daily/2026-09-15.md | 44 +++++++++++++++++++ src/shinro/codegen/compose.py | 66 +++++++++++++++++++++++++++- tests/test_codegen_build.py | 10 ++++- tests/test_codegen_compose.py | 82 +++++++++++++++++++++++++++++++++++ 5 files changed, 214 insertions(+), 8 deletions(-) diff --git a/docs/codegen.md b/docs/codegen.md index cc2a9a8..ed3d480 100644 --- a/docs/codegen.md +++ b/docs/codegen.md @@ -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,8 +185,9 @@ 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) diff --git a/lab-notes/daily/2026-09-15.md b/lab-notes/daily/2026-09-15.md index c2c703a..0c44874 100644 --- a/lab-notes/daily/2026-09-15.md +++ b/lab-notes/daily/2026-09-15.md @@ -455,4 +455,48 @@ 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/src/shinro/codegen/compose.py b/src/shinro/codegen/compose.py index e782c11..f542bd5 100644 --- a/src/shinro/codegen/compose.py +++ b/src/shinro/codegen/compose.py @@ -98,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. @@ -121,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,)) @@ -170,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 @@ -273,6 +285,7 @@ def compose( 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"]] @@ -359,7 +372,7 @@ def compose( return ComposedGraph( graph=combined, inputs=["y", "x_ref", "u_prev"] + state_ports + ctrl_state_ports + free_input_names, - outputs=["u"], + 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"], ) @@ -385,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/tests/test_codegen_build.py b/tests/test_codegen_build.py index ec8660f..a5074dc 100644 --- a/tests/test_codegen_build.py +++ b/tests/test_codegen_build.py @@ -17,7 +17,7 @@ import numpy as np import pytest -from shinro.codegen import build_composed_graph +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 @@ -325,7 +325,13 @@ def test_mppi_composes_with_a_host_filled_epsilon_port(): 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) - assert cg.outputs == ["u"] + # 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? ──────── diff --git a/tests/test_codegen_compose.py b/tests/test_codegen_compose.py index e250549..a762078 100644 --- a/tests/test_codegen_compose.py +++ b/tests/test_codegen_compose.py @@ -680,6 +680,88 @@ def test_host_port_is_a_real_graph_input(self): 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 ────────────────────────────────────────── From 4e29495a71e1943ef072c3015a7e514912e27dfe Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Wed, 16 Sep 2026 12:56:28 -0400 Subject: [PATCH 17/21] chore: added smc demo --- demos/demo_smc.py | 318 ++++++++++++++++++++++++++++++++++ lab-notes/daily/2026-09-16.md | 78 +++++++++ 2 files changed, 396 insertions(+) create mode 100644 demos/demo_smc.py create mode 100644 lab-notes/daily/2026-09-16.md 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/lab-notes/daily/2026-09-16.md b/lab-notes/daily/2026-09-16.md new file mode 100644 index 0000000..cb95681 --- /dev/null +++ b/lab-notes/daily/2026-09-16.md @@ -0,0 +1,78 @@ +# 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 + + From 4a14d5fd4ce8859b05b472511add0f56696a054b Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Wed, 16 Sep 2026 14:44:23 -0400 Subject: [PATCH 18/21] feat: make plant dynamics batch-capable for lowered nonlinear MPPI rollouts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MPPI's lowered rollout was LTI-only. `dynamics_fn` was a single batched matmul, so the graph's node count was independent of the sample count, but a nonlinear plant fell back to a per-sample Python loop that cannot be traced: each sample emits its own copy of the body, so a graph would grow as N*K (~30,000 nodes at the shipped N=200/K=15 — roughly 100 minutes of ReleaseFast compile at the measured cost per node). Make `Plant.dynamics` batch-capable, so one implementation serves the eager per-sample rollout, the finite-difference linearization, and the lowered graph: - `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. Every backend call must go through `bk`, which defaults to `self.bk`: tracing swaps only the traced component's backend, never the plant's. - New `utils/batching.py` rank helpers (`as_batch`, `as_vector`, `column`, `control_batch`) keep each physics body branch-free and in the column idiom. - All three nonlinear plants converted: InvertedPendulum, CartPole and DoublePendulum. The last needed its 2x2 solve written in closed form (Cramer's rule) — a batched mass matrix would be rank-3, which the 2-D graph backend does not represent, and the old matrix helpers used item assignment. Its determinant is strictly positive for positive masses. InvertedPendulum and CartPole `step()` now call `dynamics`, removing a third copy of their equations. - `BatchedDynamicsAdapter` collapses to "nonlinear" / "lti"; the per-sample loop, `torch.vmap`, and the intermediate `batched_dynamics` method are gone. `mppi.attach_plant` routes the controller's current backend into the plant, so a traced call emits nodes instead of evaluating eagerly. There is deliberately one implementation per plant, not two transcriptions of the same physics — the graph is the Python execution, transcribed. The physics stays pinned by the analytic tests in tests/test_plants.py (DoublePendulum Coriolis against a hand-built M and np.linalg.solve, balancing checks), and the new TestBatchCapableDynamics guards the rank contract so a future nonlinear plant cannot ship scalar-only. Verified: make test 1174 passed / 5 skipped; make lint 0 errors; generated runtime artifacts untouched. --- README.md | 4 + lab-notes/daily/2026-09-16.md | 193 +++++++++++++++++++++++++ src/shinro/components.py | 25 +++- src/shinro/controllers/mppi.py | 7 +- src/shinro/plants/cartpole.py | 46 +++--- src/shinro/plants/double_pendulum.py | 129 +++++++++-------- src/shinro/plants/inverted_pendulum.py | 36 +++-- src/shinro/utils/batched_adapter.py | 116 ++++++++------- src/shinro/utils/batching.py | 102 +++++++++++++ tests/test_batched_adapter.py | 113 +++++++++++---- tests/test_plants.py | 48 ++++++ 11 files changed, 634 insertions(+), 185 deletions(-) create mode 100644 src/shinro/utils/batching.py diff --git a/README.md b/README.md index 31f6ac0..7af7486 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,9 @@ 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_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/lab-notes/daily/2026-09-16.md b/lab-notes/daily/2026-09-16.md index cb95681..5d4041b 100644 --- a/lab-notes/daily/2026-09-16.md +++ b/lab-notes/daily/2026-09-16.md @@ -75,4 +75,197 @@ 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 + 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/controllers/mppi.py b/src/shinro/controllers/mppi.py index 59aabda..9e4678e 100644 --- a/src/shinro/controllers/mppi.py +++ b/src/shinro/controllers/mppi.py @@ -249,7 +249,12 @@ 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 + # 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, epsilon: Any | None = None): 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/utils/batched_adapter.py b/src/shinro/utils/batched_adapter.py index c7492aa..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. 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/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_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 From 0096087aeef4eac4f8c2ab36e03167555183ab6f Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Wed, 16 Sep 2026 15:45:59 -0400 Subject: [PATCH 19/21] feat: verify nonlinear MPPI end to end (interpreter oracle, .so, compile scenario) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batch-capable `Plant.dynamics` contract now has the same proof the LTI rollout has: an interpreter oracle, a compiled-kernel oracle, and an e2e scenario build. Interpreter oracle (`TestMppiNonlinearOracle`): `interpret()` vs live numpy over seeded draws (out and state_u to 1e-11), the `costs` diagnostic port, a recurrence check (the graph's own state_u feeds the next tick), a drift guard (`sin` is in the op set, `state_outputs == ["state_u"]` — state detection finds only the nominal plan, so the tracking reference is not promoted to a recurrent port), and the structural claim: tracing the same policy at N=6 and N=24 gives the identical node count while the `epsilon` port shape differs. The batch lives in the node shapes, not in the graph topology. Zig oracle (`mppi_pendulum_so`, tmp graph_path so the shipped runtime graph is never touched): three-way parity over 10 seeded draws — .so vs interpreter at 1e-13 (including costs), .so vs live numpy at 1e-11 — plus C-ABI recurrence, where the host feeds the kernel's own `state_out` back as the next `state_u`. Compile scenario (`mppi_pendulum_compile.toml` + controller config `mppi_inverted_pendulum.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`. `make compile` reports 794 nodes, outputs ['u', 'costs'], oracle .so vs interpret 2.132e-14, and a verified deployment record — the same pipeline, oracle, and stamp as the LTI path, with no special-casing. Verified: make test 1181 passed / 5 skipped; make lint 0 errors; generated runtime artifacts untouched. --- lab-notes/daily/2026-09-16.md | 159 ++++++++++++ .../controllers/mppi_inverted_pendulum.toml | 20 ++ .../scenarios/mppi_pendulum_compile.toml | 35 +++ tests/test_zig_lowering.py | 232 ++++++++++++++++++ 4 files changed, 446 insertions(+) create mode 100644 src/shinro/configs/controllers/mppi_inverted_pendulum.toml create mode 100644 tests/integration/scenarios/mppi_pendulum_compile.toml diff --git a/lab-notes/daily/2026-09-16.md b/lab-notes/daily/2026-09-16.md index 5d4041b..6e034af 100644 --- a/lab-notes/daily/2026-09-16.md +++ b/lab-notes/daily/2026-09-16.md @@ -252,6 +252,7 @@ plant cannot ship scalar-only. In `tests/test_batched_adapter.py`, the 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. @@ -268,4 +269,162 @@ 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 + 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/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_zig_lowering.py b/tests/test_zig_lowering.py index 3cb58ba..6b9b204 100644 --- a/tests/test_zig_lowering.py +++ b/tests/test_zig_lowering.py @@ -407,6 +407,65 @@ def _build_mppi_graph(N: int = MPPI_N, K: int = MPPI_K, dt: float = 0.02) -> Com ) +# 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.""" @@ -508,6 +567,20 @@ def mppi_so(tmp_path_factory): 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. @@ -1552,6 +1625,165 @@ def test_cabi_recurrence_matches_numpy(self, mppi_so): 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: """The .solve_qp VM op (codegen static solver) matches the interpreter. From 027d22404856a4a622db96323e5f14984bb1dac2 Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Wed, 16 Sep 2026 15:47:01 -0400 Subject: [PATCH 20/21] docs: document the nonlinear MPPI rollout contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MPPI's rollout is the one place a plant model enters the graph, so it is what decides whether a nonlinear plant can lower at all. `docs/codegen.md` gains a subsection under the composition pass stating the contract: `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, so one implementation serves the eager per-sample rollout, the finite-difference linearization, and the traced graph. It also states the two mechanical requirements (route every backend call through the passed `bk`, because `trace_node` swaps only the traced component's backend; and no scalar indexing, because the tracer has no `__getitem__`), the node-count argument (the batch lives in the node shapes, so a graph carries one `sin`/`mul` node of shape `(N, 1)` per term rather than N copies of the body), and the non-goal (data-dependent per-row branching belongs on the host). Also adds `min` to both op lists — it has been missing since the op landed. `demo_mppi.py`'s rule 6 no longer claims "nonlinear rollouts stay eager-only"; it now says the plant's dynamics must be batch-capable, and notes that `attach_plant` already satisfies that for both LTI (one matmul) and nonlinear (the plant's own dynamics). --- demos/demo_mppi.py | 5 ++-- docs/codegen.md | 43 +++++++++++++++++++++++++++++++---- lab-notes/daily/2026-09-16.md | 4 ++++ 3 files changed, 46 insertions(+), 6 deletions(-) diff --git a/demos/demo_mppi.py b/demos/demo_mppi.py index 049138c..ec21c3f 100644 --- a/demos/demo_mppi.py +++ b/demos/demo_mppi.py @@ -219,12 +219,13 @@ def print_trace_safety_rules(): "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", - "dynamics must be LTI for lowering; nonlinear rollouts stay eager-only", + "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 LTI path already follows all of these.") + 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(): diff --git a/docs/codegen.md b/docs/codegen.md index ed3d480..ca851a9 100644 --- a/docs/codegen.md +++ b/docs/codegen.md @@ -229,6 +229,41 @@ 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 @@ -257,8 +292,8 @@ listing available ops. The current set (from `ops.py`): `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`, `argmax`, `one_hot`, -`slice`, `solve_qp`. +`tanh`, `relu`, `exp`, `abs`, `sign`, `pow`, `sin`, `cos`, `min`, `argmax`, +`one_hot`, `slice`, `solve_qp`. ## Lowering to Zig (shipped) @@ -396,8 +431,8 @@ enum in `graph_data.zig`; `cst`/`inp`/`out`/`where_op` are the Zig spellings of `const`, `input`, `output`, `matmul`, `add`, `sub`, `mul`, `div`, `ne`, `lt`, `neg`, `transpose`, `inv`, `reshape`, `clip`, `where`, `any`, `copy`, `tanh`, -`relu`, `exp`, `abs`, `sign`, `pow`, `argmax`, `one_hot`, `slice`, `sin`, `cos`, -`stack`, `solve_qp`. +`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 interpreter handler solves with the Python `osqp` (eps=1e-6), while the VM diff --git a/lab-notes/daily/2026-09-16.md b/lab-notes/daily/2026-09-16.md index 6e034af..c9c0316 100644 --- a/lab-notes/daily/2026-09-16.md +++ b/lab-notes/daily/2026-09-16.md @@ -428,3 +428,7 @@ temp path). README demo list gained the two invocations. ### 2026-09-16 19:45 UTC — update + +### 2026-09-16 19:47 UTC — update + + From 3f8dde62f130817791df2ea93ea34c8cfb3b182a Mon Sep 17 00:00:00 2001 From: Adil Faisal Date: Wed, 16 Sep 2026 15:47:13 -0400 Subject: [PATCH 21/21] docs(demos): add a nonlinear MPPI demo, dev through compiled kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `demo_mppi.py` teaches MPPI's wiring and lowering contract on an LTI plant; this is its nonlinear twin, showing the same contract holds when the rollout is the plant's own `dynamics`. No MuJoCo or torch required. Six sections: (1) one batch-capable method, single and batch ranks agreeing; (2) the eager closed loop swinging the pendulum from 23 degrees to |theta| ~ 1.5e-3; (3) why the nonlinear rollout — the linearized acceleration is 11.5% off at 46 degrees, compounding to a 0.75 rad theta error over 0.5 s, while the plant's own dynamics is exact by construction; (4) the trace — one `sin` node per horizon step of shape (N, 1), interpreter vs live numpy at 0.0, and an identical node count at N=16 and N=64; (5) the C-ABI port table, with epsilon host-drawn; (6) `--build` lowers to a temp graph, compiles ReleaseFast, dlopens, and replays the same noise through both — max |u_kernel - u_eager| 4.55e-15 and max |x_kernel - x_eager| 2.11e-15 over 600 closed-loop ticks. Section 3 deliberately compares accelerations and uses one integrator for the horizon comparison: the plant's `step()` integrates semi-implicitly while the rollout uses explicit Euler, and mixing the two made the linearized model look better than the nonlinear one. README demo list gained the two invocations. --- README.md | 2 + demos/demo_mppi_nonlinear.py | 342 ++++++++++++++++++++++++++++++++++ lab-notes/daily/2026-09-16.md | 4 + 3 files changed, 348 insertions(+) create mode 100644 demos/demo_mppi_nonlinear.py diff --git a/README.md b/README.md index 7af7486..228a716 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,8 @@ python -m demos.demo_base_tracking # base tracking, LQR + 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 ``` 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/lab-notes/daily/2026-09-16.md b/lab-notes/daily/2026-09-16.md index c9c0316..2c684db 100644 --- a/lab-notes/daily/2026-09-16.md +++ b/lab-notes/daily/2026-09-16.md @@ -432,3 +432,7 @@ temp path). README demo list gained the two invocations. ### 2026-09-16 19:47 UTC — update + +### 2026-09-16 19:47 UTC — update + +