Skip to content

01 — Architecture

This section is the stable architectural reference for Lensemble: the module map, the dependency layering, what is federated versus local, the two-level training topology, the trust boundaries, the end-to-end data flows, the process model, and the staged rollout. It states what exists and how it fits together. Rationale for each design choice lives in the RFCs cited inline; this document does not restate rationale, it pins the contract.

Lensemble trains a single action-conditioned JEPA world model end-to-end (encoder $f_\theta$ AND predictor $g_\phi$ co-trained — "Fork B") across many mutually-distrusting participants. Raw data never leaves a participant boundary; only released model deltas and protocol metadata cross. Secure aggregation and effective DP are target protections for those deltas, not properties of the current plaintext coordinator path. The scientific core is the latent-gauge problem and its fix (the SIGReg-JEPA objective is invariant under $O(d)$ rotations of the latent space, so independently-updated participants drift into mutually-rotated frames and naive weight-averaging is meaningless). See 00 — Overview for the thesis and contribution, and RFC-0001 — Architecture & System Overview for the originating rationale.

1. Module map

The Python import root is lensemble. Every top-level module owns one subsystem, is specified by exactly one RFC, and has a defined public/internal split. The layout below is canonical (conventions §1); the reference implementation MUST NOT introduce additional top-level modules without a new RFC. Anything under a module named _internal or prefixed with _ is private and unversioned (see 02 — Public API §2 for the stability policy).

Module Responsibility Specified by Surface
lensemble/__init__.py Public re-exports and __version__ (SemVer). conventions §5 public
errors.py Error taxonomy: LensembleError base, LensembleErrorCode enum, every typed error carrying .code and .remediation. 04 — Error Model public
cli.py Typer CLI app (train, federate, eval, probe, commit, drift, verify, doctor); every command emits a RunManifest. 02 — Public API §3 public
contracts/ WMCP latent contract (LatentState, ActionSpec) and embodiment conformance checks. RFC-0007 public
model/ Encoder $f_\theta$, predictor $g_\phi$, per-embodiment action heads $h_\psi^{(c)}$, the three-term Objective, the SIGReg statistic. Files: encoder.py, predictor.py, action_head.py, objective.py, sigreg.py. RFC-0008 public
gauge/ Frame anchoring, Procrustes alignment, frame-drift diagnostics. Files: anchor.py, procrustes.py, drift.py. RFC-0002 public
federation/ DiLoCo outer loop, round state machine, roles. Files: coordinator.py, participant.py, round.py, outer_optimizer.py. RFC-0003, RFC-0013 public
aggregation/ Secure-aggregation primitives (pairwise masking / TEE), deterministic summation, and current post-commit equivalence reporting. Files: secure_agg.py, masking.py. RFC-0011 internal
privacy/ Clip+noise primitives, $(\varepsilon,\delta)$ accountant primitives, and current one-round replay reporting. Files: dp.py, accountant.py. RFC-0012 public (accountant), internal (mechanism)
data/ Data layer, window loader, embodiment adapters, residency enforcement, public probe. Files: dataset.py (the EpisodeDataset window loader), adapters/, residency.py, probe.py. RFC-0004 public
provenance/ Episode hashing, Merkle tree, contribution ledger. Files: merkle.py, commit.py, ledger.py. RFC-0014 public
eval/ Latent MPC planner, eval harness, metrics. Files: mpc.py, harness.py, metrics.py. RFC-0005 public
config/ Structured config schema, run manifest, seeding. Files: schema.py, manifest.py, seed.py. RFC-0009 public
artifacts/ Checkpoint/artifact format, hashing, schema versioning. Files: checkpoint.py, schema.py, hashing.py. RFC-0010 public
observability/ Structured logging, metric emission, redaction guard. Files: logging.py, metrics.py, redaction.py. RFC-0015 public
verify/ Phase-2 verifiable layer + Phase-1 public recomputation. Files: recompute.py, stark.py. RFC-0006 public (recompute), stub (stark)

tests/ holds the pytest suite (unit, property, integration, ml); docs/ holds this corpus; configs/ holds the Hydra config groups. These are not import modules.

The model split is fixed (conventions §2, RFC-0008 — Proposed Design):

  • Encoder $f_\theta:\text{video clip}\to\mathbb{R}^{N\times d}$ — a video Vision Transformer warm-started from released V-JEPA 2 weights, co-trained under SIGReg. The round-0 snapshot is $f_{\text{ref}}$, frozen and used as the gauge anchor (INV-WARMSTART-T0).
  • Predictor $g_\phi$ — a compact transformer predicting future latents autoregressively, conditioned on an action embedding (the LeWM ARPredictor shape).
  • Action head $h_\psi^{(c)}$ — per-participant, mapping that embodiment's action space into the shared latent-conditioning space. Never broadcast or aggregated (INV-ACTIONHEAD-LOCAL).

The shared latent interface that every encoder emits and every predictor consumes is the WMCP LatentState contract (RFC-0007). It makes heterogeneous-embodiment federation type-safe — the explicit analogue of the fixed token vocabulary that LLM federation gets for free.

2. Dependency layering

The modules form an acyclic dependency graph in four bands. A lower band never imports a higher one. core here means the always-importable foundation: errors.py, config/, and observability/ (these carry no domain dependencies and may be imported by anything). The acyclicity is enforced by a module-boundary import test (see 07 — Testing Strategy and RFC-0001 — Testing Strategy).

Band 0  core            errors.py   config/   observability/
                            ▲          ▲           ▲
Band 1  contracts        contracts/  artifacts/  provenance/   privacy/
        + leaf io           ▲          ▲           ▲              ▲
Band 2  model/gauge      model/  ─────┘            │              │
        + aggregation    gauge/  ─────┘            │              │
                         aggregation/ ─────────────┘              │
                            ▲                                     │
Band 3  orchestration    federation/ ── imports ──► model, gauge, aggregation,
                            ▲                        privacy, provenance, artifacts
                         eval/ ── imports ──► model, data
                            ▲
                         cli.py / verify/ ── top-level entry points

In prose, the allowed dependency edges are:

  • model/ and gauge/ depend on contracts/ (they produce and align LatentState instances).
  • aggregation/ depends on contracts/ and artifacts/ only for the delta and commitment types it transports; it depends on privacy/ for the clip+noise step ordering (RFC-0011).
  • federation/ (the orchestrator) depends on model/, gauge/, aggregation/, privacy/, provenance/, and artifacts/. It is the only module that composes the full round.
  • eval/ depends on model/ (to load and run the world model) and data/ (to source clips and the probe); it does NOT depend on federation/.
  • data/ depends on contracts/ (loaders emit WMCP-conformant windows) and provenance/ (commit on ingest).
  • verify/ depends on gauge/, provenance/, and artifacts/ for public recomputation; the Phase-2 stark.py path is a stub raising NotImplementedError.
  • Everything may depend on Band 0: errors.py, config/, observability/.

There are no back-edges: contracts/ imports nothing from model/, data/, or federation/; model/ imports nothing from federation/ or eval/. The orchestration band is the apex. The boundary-import test imports each module in isolation and asserts no cycle and no upward edge; a violation is a build failure, not a runtime error.

3. Federation map

Which components are federated, local, shared-per-round, or never-crossing is fixed by the model split and the gauge mechanism. This table expands RFC-0001 §4 and is the authoritative disposition; the invariants named are enforced as cited.

Component Disposition Enforced by / invariant Why
Encoder backbone $f_\theta$ Federated (gauge-controlled) INV-WARMSTART-T0 at $t{=}0$; frame anchor each round Shared physics; the point of the project
Predictor core $g_\phi$ Federated aggregated alongside $f_\theta$ in $\Delta_c$ Shared dynamics; frame-pinned so averaging is valid
Action head $h_\psi^{(c)}$ Local — personalized INV-ACTIONHEAD-LOCAL (RFC-0007) Embodiment action spaces differ (quadruped ≠ 7-DoF arm)
SIGReg sketch matrix $A$ Shared per round (broadcast seed $s_t$) INV-SKETCH-CONSISTENCY (RFC-0002 §3) Objective consistency across silos in a round
Public probe $\mathcal{P}$ + landmarks ${t_i}$ Shared, fixed, hash-pinned INV-PROBE-PIN (RFC-0004) The manufactured frame anchor
Reference encoder $f_{\text{ref}}$ Shared, frozen INV-WARMSTART-T0 Defines round-0 frame; landmark targets $t_i=f_{\text{ref}}(p_i)$
Pseudo-gradient $\Delta_c$ Crosses boundary (DP + secure-agg) INV-DP-BOUND, INV-COMMIT-BINDING The only learning signal that leaves a participant
Dataset Merkle root $R_c$ Crosses boundary INV-COMMIT-BINDING (RFC-0014) Binds a contribution to its data, not the data itself
Global params $(\theta_t,\phi_t)$ hash Shared coordination state INV-CHECKPOINT-HASH (RFC-0010) Reproducible, committed round state
Raw trajectories (obs/actions) Never leaves the boundary INV-RESIDENCY (RFC-0004) Sovereignty (fail-closed)
Private-data embeddings Never leaves the boundary INV-RESIDENCY Sovereignty (fail-closed)

INV-RESIDENCY is enforced in lensemble.data.residency: any attempt to serialize a raw observation, action, or private-data embedding into an outbound message or artifact raises ResidencyViolation, which is fail-closed and never caught-and-ignored (04 — Error Model, 06 — Security). The probe $\mathcal{P}$ is public/licensed data and its embeddings $E_{\text{ref}}=f_{\text{ref}}(\mathcal{P})$ are not private, so they may cross for alignment.

4. Training topology

Training nests two loops. The standard distributed-training stack is the inner loop and is not the contribution; the federated outer loop and its gauge control are. See RFC-0001 §5 and RFC-0003 §2.

  • Inner — intra-participant, for scale. Within a participant, standard FSDP2 / tensor / context parallelism (torch >=2.4) trains the warm-started model (ViT-L/~300M at Stage A, toward the 1.2B target). The inner optimizer is AdamW. SIGReg projection statistics are reduced within this trust domain freely — the reduce-within-trust-domain rule (RFC-0008). This is the only place the large-model-parallelism playbook applies; it is off-the-shelf.
  • Outer — inter-participant, for sovereignty. DiLoCo: each participant runs $H$ inner steps, then the coordinator applies an outer Nesterov step on the averaged pseudo-gradient $(\theta_{t+1},\phi_{t+1}) = (\theta_t,\phi_t) + \eta_{\text{out}}\,\mathrm{Nesterov}\big(\tfrac{1}{C}\sum_c\Delta_c\big)$, where $\Delta_c = (\theta_c^{\text{local}},\phi_c^{\text{local}}) - (\theta_t,\phi_t)$. On the target path, any participant-specific gauge transform is applied to raw local full weights and re-differenced first; only the resulting aligned $\Delta'_c$ crosses after clip/noise, optional quantization, and masking (INV-ALIGN-BEFORE-RELEASE). In the current runtime, each released $\Delta_c$ is visible to the coordinator; secure aggregation is not on the optimizer path. Sync frequency is every $H$ inner steps; $H\in[50,500]$ (smaller while characterizing drift). $H$ is tuned against frame drift (RFC-0002 §4): larger $H$ is cheaper communication but more drift to anchor. Lensemble adds its local-minus-global displacement; this is equivalent to subtracting DiLoCo's global-minus-local outer gradient.

The outer step is on the bitwise-deterministic aggregation path (INV-AGG-DETERMINISM, enforced in lensemble.aggregation / lensemble.federation.outer_optimizer): fixed reduction order, fp32 with fixed summation order, no atomics. A determinism self-check runs each outer step; failure raises NonDeterministicAggregation and aborts the round (conventions §9, 04 — Error Model). The inner loop is best-effort deterministic and seed-pinned; full inner determinism is gated by a config flag (torch.use_deterministic_algorithms).

5. Trust boundaries

Each participant is a sovereign trust domain. The coordinator/aggregator is a separate domain, treated as honest-but-curious in Phase 1 and as a proving target in Phase 2 (RFC-0006). This diagram reproduces RFC-0001 §6; the prose below says the same.

┌── Participant c (sovereign) ─────────────────────────┐
│  raw trajectories  ──►  local train (inner-parallel) │
│        │ (never leaves)         │                     │
│        ▼                        ▼                     │
│  Merkle commitment R_c     raw local W_c / Δ_c         │
│                            │ align + re-difference      │
│                            ▼                            │
│                         aligned Δ'_c                    │
└───────────────│──────────────────│─────────────────────┘
                │                  │  (clip + replay noise in current path)
                ▼                  ▼
        ┌──────────── Current coordinator ─────────────────────┐
        │  individual released Δ_c visible; plaintext commit    │
        │  outer Nesterov step → θ^{global}_{t+1} (hash-committed)│
        │  optional coordinator alignment: raw plaintext harness │
        └───────────────────────────────────────────────────────┘

What crosses a boundary: released model deltas $\Delta_c$, dataset commitments $R_c$, and shared coordination state (the sketch seed $s_t$, the probe hash, the global-model hash). The current coordinator sees each released $\Delta_c$ and commits from their plaintext reduction; simulated/TEE secure-sum results are post-commit equivalence checks, and masking reports explicit fallback. What never crosses: raw observations, actions, or embeddings of private data — enforced by INV-RESIDENCY. Public recompute_alignment measures alignment of a committed global model on the public probe; it does not verify a participant-specific pre-release transform. The current coordinator transform needs individually visible raw deltas and is therefore restricted to a privacy-off, quantization-off, simulated plaintext harness. A genuine sum-only path must align participant-side first or use a trusted/MPC boundary (RFC-0006). The boundary-crossing message set and per-message protection are tabulated in 06 — Security and RFC-0003 §8.

6. Data flow lifecycles

6.1 Federated training round (normative Stage-C target)

End-to-end, per round $t$, mapping to RFC-0003 §2:

  1. Coordinator broadcasts RoundOpen: the global parameters $(\theta_t,\phi_t)$ reference/hash, the round sketch seed $s_t$, the probe hash, and $H$. Action heads $h_\psi^{(c)}$ are not broadcast (INV-ACTIONHEAD-LOCAL).
  2. Each participant verifies the probe hash against the pinned probe (INV-PROBE-PIN; mismatch raises ProbeError) and derives the identical sketch matrix $A$ from $s_t$ (INV-SKETCH-CONSISTENCY).
  3. Local optimization: $H$ inner AdamW steps on the objective $\mathcal{L} = \lambda_{\text{pred}}\,\mathbb{E}\lVert g_\phi(f_\theta(x_t),a_t) - \text{sg}[f_\theta(x_{t+1})]\rVert^2 + \lambda_{\text{sig}}\,\mathrm{SIGReg}A(f\theta(x)) + \lambda_{\text{anc}}\,\mathcal{L}{\text{anchor}}(f\theta;\mathcal{P},{t_i})$, over local data only.
  4. Form the raw displacement $\Delta_c = (\theta_c^{\text{local}},\phi_c^{\text{local}}) - (\theta_t,\phi_t)$.
  5. Backstop align before release: if frame drift exceeds the threshold, compute $Q_c^\star$ on the public probe, transform raw local full weights, and re-difference $\Delta'_c=W'_c-W_t$; otherwise $\Delta'_c=\Delta_c$ (INV-ALIGN-BEFORE-RELEASE; RFC-0002 §5).
  6. Privatize: clip $\Delta'c \leftarrow \Delta'_c\cdot\min(1, C{\text{clip}}/\lVert\Delta'c\rVert)$ (INV-DP-BOUND, post-clip $\lVert\Delta'_c\rVert \le C{\text{clip}}$), then add $\mathcal{N}(0,\sigma^2 C_{\text{clip}}^2 I)$ (RFC-0012). The target accountant updates cumulative $(\varepsilon,\delta)$ before release; exhaustion raises PrivacyBudgetExceeded and stops training. The current report instead constructs a fresh one-round accountant after commit and does not enforce a cumulative budget.
  7. Each participant optionally quantizes, encodes, and masks the aligned private release, then emits Commitment carrying its dataset Merkle root $R_c$, binding $\Delta'_c$ to exactly one $R_c$ (INV-COMMIT-BINDING; mismatch raises CommitmentMismatch, update rejected).
  8. Secure-aggregate: compute $\sum_c\mathrm{released}(\Delta'_c)$ without revealing any individual update (RFC-0011). Dropout below threshold raises SecureAggregationError; participant churn below the minimum raises FaultToleranceExceeded. This is the target path. The current optimizer consumes plaintext individual deltas before any simulated/TEE equivalence report is built.
  9. Outer step: deterministic Nesterov update producing $(\theta_{t+1},\phi_{t+1})$ (INV-AGG-DETERMINISM; self-check failure raises NonDeterministicAggregation, aborts and recomputes).
  10. Commit: hash-commit $(\theta_{t+1},\phi_{t+1})$ (INV-CHECKPOINT-HASH); emit RoundClose with the hash; append a ContributionRecord to the ledger (RFC-0014).
RoundOpen → [verify probe + derive A] → H inner steps → raw W_c/Δ_c
  → (drift? align W_c; re-difference Δ'_c) → clip+noise → optional quantize → encode+mask
  → Commitment(R_c) → secure-agg Σ_c released(Δ'_c) → Nesterov outer step
  → hash-commit θ_{t+1} → RoundClose → ledger append

6.2 Single-site Stage-A training

  1. Load LensembleConfig (Hydra structured config); validate (raises ConfigError); emit a RunManifest (config hash, seeds, env, versions, git SHA).
  2. Build the WMCP-conformant encoder (warm-started from V-JEPA 2), predictor, and a single action head; snapshot $f_{\text{ref}}$.
  3. Train end-to-end on pooled robot data with the three-term objective. This is the centralized upper bound (no federation, no DP, no secure aggregation). Public entry point train_local(config) -> RunResult (02 — Public API).
  4. Periodically checkpoint via the schema-versioned, hash-committed artifact format (RFC-0010).
  5. Evaluate via latent MPC (§6.3) to validate the objective before federation.

6.3 Evaluation (latent MPC)

  1. Load a committed checkpoint (verify INV-CHECKPOINT-HASH; tamper raises CheckpointIntegrityError).
  2. Encode the goal and current observation to latent states via $f_\theta$.
  3. The Planner runs latent model-predictive control — CEM / iCEM / MPPI minimizing an $L_1$ goal-energy in latent space, rolling $g_\phi$ forward under candidate action sequences, exactly as stable-worldmodel provides (RFC-0005).
  4. The eval harness records eval/success_rate, eval/planning_samples, eval/time_per_action_ms and produces an EvalReport. Public entry point evaluate(checkpoint, env_id, *, cfg) -> EvalReport.

6.4 Dataset commit

  1. On ingest, each episode is canonically serialized and hashed to a domain-separated SHA-256 leaf.
  2. Leaves are assembled into a Merkle tree with defined ordering; the root is $R_c$ (RFC-0014).
  3. commit_dataset(dataset) -> DatasetCommitment returns the root $R_c$, episode count, and WMCP metadata. The commitment binds future pseudo-gradients to this data (INV-COMMIT-BINDING); a verification failure raises MerkleVerificationError.
  4. This runs from day one for Phase-1 tamper-evidence and is the bridge to Phase-2 proofs (RFC-0006).

7. Process and concurrency model

  • Coordinator process. One per federation. Runs the round state machine OPEN → COLLECTING → AGGREGATING → ALIGNING → COMMITTING → CLOSED (with an ABORTED path), holds the canonical global model, runs the outer optimizer on the deterministic path, and emits the contribution ledger (RFC-0013). Public entry point Coordinator.run(num_rounds: int) -> None.
  • Participant processes. One per sovereign node. Each runs Participant.local_round(global_state, round_seed) -> PseudoGradient, performing the $H$ inner steps and emitting a transformed, commitment-bound delta. The current clip/replay-noise transform is not effective DP. Backpressure and timeouts on the collection step are handled by the coordinator; a participant that misses the window is reconciled next round (elasticity, RFC-0013).
  • Inner parallel workers. Within a participant, FSDP2 / tensor / context parallelism distributes the large model across local devices. This is intra-trust-domain; no Lensemble boundary is crossed here.
  • Simulation harness. In Stage B the coordinator and participants run in one process (in-process round simulation) so the gauge science is studied on one cluster; in Stage C they run as separate networked processes over a real boundary (RFC-0013). Both stages drive the SAME RoundState machine and the SAME operation-oriented Transport seam (lensemble.federation.transport.Transport); only the transport changes. Stage C's NetworkedTransport (lensemble.federation.network) implements that seam over the low-level RFC-0013 §5 MessageChannel wire layer (send / recv with the None-on-timeout contract / broadcast / peers) carrying the four ControlMessages (RoundOpen / Commitment / Update / RoundClose, lensemble.federation.messages). It is interchangeable with InProcessTransport in Coordinator(config, *, transport=...): a two-node round over the in-process LoopbackChannel commits a global_state_hash() bit-identical to the in-process run on the same seed (the deterministic outer step makes the equality exact). The four messages carry only hashes, coordination scalars, and the released Δ_c — never raw data (INV-RESIDENCY, fail-closed). The current network path exposes that delta to the coordinator; masked ingress remains a target contract. Ingress re-validates every message and binds each Δ_c to its committed R_c (CommitmentMismatch, INV-COMMIT-BINDING, never swallowed). The concrete wire protocol (gRPC vs HTTP/REST) is bound at Stage C behind the MessageChannel Protocol — the state machine is transport-agnostic (RFC-0013 §Open Questions).

Device contract (conventions §9): CUDA primary, with a CPU fallback path that runs the small CI configs; tests must pass on CPU. Compute dtype is bf16 forward with fp32 master weights and loss/statistic accumulation; the aggregation path is fp32 with fixed summation order (or fp64).

8. Target staged plan mapped to milestones

The A–E staging from RFC-0001 — Migration / Rollout maps target capabilities onto the release milestones (conventions §12). The table is a roadmap, not a statement that the corresponding scientific result or trust property has been achieved.

Stage Goal Milestone
A Single-site, warm-started, ViT-L/~300M end-to-end SIGReg + AC predictor on pooled robot data; latent-MPC eval. The centralized upper bound, plus foundational scaffolding (package skeleton, config, data layer, WMCP contract, model+objective, eval harness, observability, artifact format, error taxonomy, CI, packaging). v0.1
B Simulated federation on one cluster: DiLoCo outer loop, frame anchor (Layers 1–4), Procrustes backstop, aggregation/DP mechanism plumbing with explicit effectiveness status, the frame-drift diagnostic, the full ablation ladder and non-IID/scale sweeps. The scientific core / the paper. v0.2
C Two real sovereign nodes over a network boundary: real secure aggregation + DP, residency enforcement, fault tolerance/elasticity, contribution ledger. The sovereignty demonstration. v0.3
(hardening) Frozen public API, complete docs + reproducibility package, release automation, Fork A fallback supported and tested, and deterministic/provenance disciplines verified end-to-end. v1.0
D The actual STARK/TEE verifiable layer (Phase 2) beyond the proof-ready disciplines. Out of v1.0 scope (future work)
E Own foundation-scale federated video pretraining (toward 1.2B from scratch). Out of v1.0 scope (future work)

Proof-ready engineering disciplines (deterministic aggregation, hash commitments, Merkle roots, pinned probe, public recomputation) are in scope for v0.1–v1.0; the actual proofs (Stage D) and own-pretrain (Stage E) are future work captured in the tracker, not implementable v1.0 issues (conventions §12, RFC-0006).

Fork B (end-to-end, encoder co-trained) is the target architecture. Fork A (frozen shared encoder, federate the predictor only) is the documented safe-degrade fallback: it dissolves the latent gauge entirely while narrowing the research scope to predictor federation, and it must be supported and tested by v1.0 (RFC-0002, RFC-0008).

9. Open questions

OPEN QUESTION: The personalization boundary — how heterogeneous the participating embodiments can be before a single global encoder $f_\theta$ stops being viable and the shared-backbone + per-embodiment-head split breaks. The split is the hedge; the boundary is empirical. Owner @AbdelStark; resolution path: the non-IID severity sweep in Stage B (RFC-0001, RFC-0005).

RISK: The architecture depends on a released V-JEPA 2 warm-start for both initialization and the $t{=}0$ frame anchor (INV-WARMSTART-T0). If the warm-start is unavailable or unsuitable at the target scale, both the gauge closure and the modest-compute staging assumption weaken. Resolution plan: Stage A validates the objective and eval on the warm-start before any federation work; if warm-start quality is insufficient, the Fork A fallback (frozen encoder) still recovers a clean federation (RFC-0002).

10. References