nonlinear controllers (SMC, MPPI), batched dynamics and operations - #22
Merged
Merged
Conversation
adilfaisal01
changed the title
… stack, solve_qp) with Zig lowering updates and documentation enhancements
… and update documentation
… comprehensive tests
…pe in documentation
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.
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.
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.
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.
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.
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.
`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.
…bservable 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.
…llouts 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.
…ile scenario) 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.
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).
`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.
adilfaisal01
force-pushed
the
feat/nonlinear-zig-lowering
branch
from
September 16, 2026 20:55
2787c18 to
3f8dde6
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.