v2

model.py

@@ -0,0 +1,509 @@
+"""
+model.py — Degradation Function Estimation
+
+Implements the model from the project PDF:
+
+    dhi/dt = K * (∏_c I_c) * ∑_j (θ_ij,I * h_j + θ_ij,II * h_j * ln(h_j))
+
+K is absorbed into theta_I and theta_II, so no separate input-scaling parameter is estimated.
+
+Parameters:
+    theta_I     : (N, N) — linear health coupling; theta_I[i,j] = θ_ij,I
+    theta_II    : (N, N) — log-linear health coupling; theta_II[i,j] = θ_ij,II
+
+Adjacency mask A (N x N, binary):
+    A[i, j] = 1  iff component j is physically allowed to influence component i.
+    Built from component_graph.COMPONENT_GRAPH so that the learned theta matrices
+    can only be non-zero where a real physical coupling exists.  The mask is applied
+    element-wise:   theta_I_eff = A * theta_I
+    Positions where A[i,j] = 0 are zeroed on init and their gradients are zeroed
+    during fitting — the model cannot learn phantom interactions.
+
+Stochastic extension (§5): each Euler step subtracts Q_i * H_i where
+    Q_i ~ N(0,1)²  (squared standard normal — event intensity, always ≥ 0)
+    H_i ~ Poisson(λ_i * τ / n)  (event count per step; λ_i set per component)
+Negative health values are valid in this mode and represent catastrophic failure.
+
+Fitting: Euler-forward simulation + manual Jacobian recurrence → SGD + L1 regularisation.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+
+# ---------------------------------------------------------------------------
+# Data structures
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Sample:
+    """One training example."""
+    X: np.ndarray   # (C,) constant input vector over [0, tau]
+    tau: float      # time horizon
+    y: np.ndarray   # (N,) target health at tau
+    M: np.ndarray   # (N,) mask: 1 = observed, 0 = ignored
+
+    @classmethod
+    def from_json(cls, path: str) -> List["Sample"]:
+        """
+        Load a list of samples from a JSON file.
+
+        Expected format — a JSON array where each element has:
+            "X"   : list of C floats  (inputs)
+            "tau" : float             (time horizon)
+            "y"   : list of N floats  (target health per component)
+            "M"   : list of N ints    (mask: 1 = observed, 0 = ignored)
+        """
+        with open(path, "r") as f:
+            records = json.load(f)
+        return [
+            cls(
+                X   = np.array(r["X"],   dtype=float),
+                tau = float(r["tau"]),
+                y   = np.array(r["y"],   dtype=float),
+                M   = np.array(r["M"],   dtype=float),
+            )
+            for r in records
+        ]
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+class DegradationModel:
+    """
+    Parametric ODE model for multi-component health degradation.
+
+    State starts at h(0) = 1 (all components fully healthy).
+    The ODE for component i couples to every component j via theta_I and theta_II,
+    and to every input c via theta_input.
+    """
+
+    def __init__(
+        self,
+        N: int,
+        C: int,
+        lambda_rates: Optional[np.ndarray] = None,
+        adjacency_mask: Optional[np.ndarray] = None,
+        seed: int = 0,
+    ) -> None:
+        self.N = N
+        self.C = C
+        # Per-component Poisson rates λ_i for the stochastic shock term (§5).
+        # Defaults to zeros — no randomness unless explicitly set.
+        self.lambda_rates: np.ndarray = (
+            np.asarray(lambda_rates, dtype=float)
+            if lambda_rates is not None
+            else np.zeros(N, dtype=float)
+        )
+
+        # Adjacency mask A[i,j] = 1 iff component j may influence component i.
+        # If None, defaults to all-ones (fully connected — no structural constraint).
+        # Pass build_adjacency_matrix() to enforce the physical graph topology.
+        if adjacency_mask is not None:
+            self.A: np.ndarray = np.asarray(adjacency_mask, dtype=float)
+            if self.A.shape != (N, N):
+                raise ValueError(
+                    f"adjacency_mask must be ({N},{N}), got {self.A.shape}"
+                )
+        else:
+            self.A = np.ones((N, N), dtype=float)
+
+        rng = np.random.default_rng(seed)
+        # Small negative diagonal drives self-degradation; off-diagonals start near 0.
+        # Mask is applied immediately so forbidden positions start at exactly 0.
+        self.theta_I: np.ndarray = -np.abs(rng.normal(0.0, 1e-3, (N, N))) * self.A
+        np.fill_diagonal(self.theta_I, -1e-2)   # diagonal always in mask (self-coupling)
+        self.theta_II: np.ndarray = np.zeros((N, N), dtype=float)
+
+    # ------------------------------------------------------------------
+    # Parameter vector helpers
+    # ------------------------------------------------------------------
+
+    @property
+    def num_params(self) -> int:
+        return 2 * self.N * self.N
+
+    def _get_params(self) -> np.ndarray:
+        return np.concatenate([
+            self.theta_I.ravel(),
+            self.theta_II.ravel(),
+        ])
+
+    def _set_params(self, p: np.ndarray) -> None:
+        N = self.N
+        self.theta_I  = np.minimum(p[: N * N].reshape(N, N), 0.0)
+        self.theta_II = np.maximum(p[N * N :].reshape(N, N), 0.0)
+
+    # ------------------------------------------------------------------
+    # ODE
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _safe_hlog(h: np.ndarray) -> np.ndarray:
+        """h * ln(h) with h clipped away from 0."""
+        h_safe = np.clip(h, 1e-10, None)
+        return h_safe * np.log(h_safe)
+
+    def _P(self, I: np.ndarray) -> float:
+        """Input product P = ∏_c I_c."""
+        return float(np.prod(I))
+
+    def f(self, h: np.ndarray, I: np.ndarray) -> np.ndarray:
+        """Rate vector dh/dt, shape (N,).
+
+        The mask A is applied element-wise before the matrix products so that
+        forbidden couplings (A[i,j]=0) never contribute to dh/dt regardless of
+        the current value of theta_I or theta_II.
+        """
+        P     = self._P(I)
+        h_log = self._safe_hlog(h)
+        g = (self.A * self.theta_I) @ h + (self.A * self.theta_II) @ h_log  # (N,)
+        return P * g
+
+    # ------------------------------------------------------------------
+    # Forward simulation (Euler integration)
+    # ------------------------------------------------------------------
+
+    def simulate(
+        self,
+        X: np.ndarray,
+        tau: float,
+        n_steps: int = 100,
+        stochastic: bool = False,
+        seed: Optional[int] = None,
+    ) -> np.ndarray:
+        """
+        Integrate from h(0)=1 to h(tau) using Euler steps.
+
+        When stochastic=True, each step subtracts a random shock Q_i * H_i where
+            Q_i ~ N(0,1)²   (squared standard normal — event intensity)
+            H_i ~ Poisson(lambda_rates[i] * tau / n_steps)
+        Negative health values are kept as-is; they represent catastrophic failure.
+
+        Returns shape (n_steps + 1, N) — row 0 is h(0), row k is h(k * tau / n_steps).
+        """
+        h = np.ones(self.N, dtype=float)
+        dt = tau / n_steps
+        rng = np.random.default_rng(seed)
+        trajectory = [h.copy()]
+
+        for _ in range(n_steps):
+            dh = dt * self.f(h, X)
+            if stochastic:
+                Q = rng.standard_normal(self.N) ** 2                    # N(0,1)²
+                H = rng.poisson(self.lambda_rates * tau / n_steps)      # Poisson(λi*τ/n)
+                dh -= Q * H
+            h = h + dh
+            if not stochastic:
+                h = np.clip(h, 0.0, 1.0)
+            trajectory.append(h.copy())
+
+        return np.array(trajectory)  # (n_steps + 1, N)
+
+    # ------------------------------------------------------------------
+    # Loss
+    # ------------------------------------------------------------------
+
+    def loss(self, y_hat: np.ndarray, y: np.ndarray, M: np.ndarray) -> float:
+        """Masked MSE: ∑_i (ŷi - yi)² * Mi."""
+        return float(np.sum((y_hat - y) ** 2 * M))
+
+    # ------------------------------------------------------------------
+    # Jacobians
+    # ------------------------------------------------------------------
+
+    def _df_dtheta(self, h: np.ndarray, I: np.ndarray) -> np.ndarray:
+        """
+        ∂f/∂θ, shape (N, num_params).
+
+        Columns correspond to [theta_I (flattened) | theta_II (flattened)].
+        Gradient columns for positions where A[i,j]=0 are zeroed so those
+        parameters receive no update signal during backprop.
+        """
+        N = self.N
+        P = self._P(I)
+        h_log = self._safe_hlog(h)
+
+        jac = np.zeros((N, self.num_params))
+
+        # ∂fi/∂θij,I = P * hj  — then zero out forbidden positions via A
+        raw_I = np.kron(np.eye(N), P * h[np.newaxis, :])       # (N, N*N)
+        jac[:, : N * N] = raw_I * self.A.ravel()[np.newaxis, :]
+
+        # ∂fi/∂θij,II = P * hj * ln(hj) — same masking
+        raw_II = np.kron(np.eye(N), P * h_log[np.newaxis, :])  # (N, N*N)
+        jac[:, N * N :] = raw_II * self.A.ravel()[np.newaxis, :]
+
+        return jac
+
+    def _df_dh(self, h: np.ndarray, I: np.ndarray) -> np.ndarray:
+        """
+        ∂f/∂h (Jacobian of rate w.r.t. state), shape (N, N).
+
+        ∂fi/∂hj = KP * (θij,I + θij,II * (1 + ln(hj)))
+        """
+        P = self._P(I)
+        h_safe = np.clip(h, 1e-10, None)
+        d_log = 1.0 + np.log(h_safe)   # d/dhj [hj ln hj] = 1 + ln hj
+        return P * (self.theta_I + self.theta_II * d_log[np.newaxis, :])
+
+    # ------------------------------------------------------------------
+    # Gradient via Jacobian recurrence
+    # ------------------------------------------------------------------
+
+    def compute_gradient(
+        self,
+        sample: Sample,
+        n_steps: int = 50,
+        J_clip: float = 1e6,
+    ) -> Tuple[np.ndarray, float]:
+        """
+        Gradient of the masked MSE loss for one sample, via:
+
+            Jθŷ(t + dt) = Jθŷ(t) + dt * (∂f/∂θ + ∂f/∂ŷ · Jθŷ(t))
+            Jθŷ(0) = 0
+
+        where Jθŷ = ∂ŷ/∂θ has shape (N, num_params).
+
+        Returns (gradient w.r.t. params, scalar loss).
+        """
+        X, tau, y, M = sample.X, sample.tau, sample.y, sample.M
+        dt = tau / n_steps
+        N, P = self.N, self.num_params
+
+        h = np.ones(N, dtype=float)
+        J = np.zeros((N, P))    # Jθŷ
+
+        for _ in range(n_steps):
+            df_dt = self._df_dtheta(h, X)  # (N, P)
+            df_dh = self._df_dh(h, X)       # (N, N)
+            J = np.clip(J + dt * (df_dt + df_dh @ J), -J_clip, J_clip)
+            h = np.clip(h + dt * self.f(h, X), 0.0, 1.0)
+
+        # ∂L/∂θr = 2 * ∑_i (ŷi − yi) * Mi * ∂ŷi/∂θr
+        residual = (h - y) * M      # (N,)
+        grad = 2.0 * (residual @ J) # (P,)
+        loss_val = float(np.sum(residual ** 2))
+        return grad, loss_val
+
+    # ------------------------------------------------------------------
+    # Summary
+    # ------------------------------------------------------------------
+
+    def summary(self, component_names: Optional[List[str]] = None) -> None:
+        """Print a human-readable overview of the fitted parameters."""
+        comp = component_names or [f"comp_{i}" for i in range(self.N)]
+        w = max(len(n) for n in comp)   # column width
+
+        print(f"DegradationModel  N={self.N}  C={self.C}  params={self.num_params}")
+        print()
+
+        print("Linear health coupling (theta_I[i,j])  -- row i influenced by col j:")
+        header = " " * (w + 4) + "  ".join(f"{n:>{w}}" for n in comp)
+        print(header)
+        for i, row_name in enumerate(comp):
+            vals = "  ".join(
+                f"{self.theta_I[i, j]:+{w}.4f}" if self.A[i, j] else " " * (w + 1) + "-"
+                for j in range(self.N)
+            )
+            print(f"  {row_name:<{w}}  {vals}")
+        print()
+
+        print("Log-linear health coupling (theta_II[i,j])  -- row i influenced by col j:")
+        print(header)
+        for i, row_name in enumerate(comp):
+            vals = "  ".join(
+                f"{self.theta_II[i, j]:+{w}.4f}" if self.A[i, j] else " " * (w + 1) + "-"
+                for j in range(self.N)
+            )
+            print(f"  {row_name:<{w}}  {vals}")
+        print()
+
+        print("Stochastic shock rates (lambda_rates):")
+        for i, (name, lam) in enumerate(zip(comp, self.lambda_rates)):
+            print(f"  {name}: lambda={lam:.6f}")
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def save(self, path: str) -> None:
+        """Save all model arrays to a .npz file."""
+        np.savez(
+            path,
+            N            = self.N,
+            C            = self.C,
+            lambda_rates = self.lambda_rates,
+            A            = self.A,
+            theta_I      = self.theta_I,
+            theta_II     = self.theta_II,
+        )
+
+    @classmethod
+    def load(cls, path: str) -> "DegradationModel":
+        """Load a model saved with save()."""
+        d = np.load(path)
+        m = cls(
+            N              = int(d["N"]),
+            C              = int(d["C"]),
+            lambda_rates   = d["lambda_rates"],
+            adjacency_mask = d["A"],
+        )
+        m.theta_I  = np.minimum(d["theta_I"],  0.0)
+        m.theta_II = np.maximum(d["theta_II"], 0.0)
+        return m
+
+    # ------------------------------------------------------------------
+    # Fitting (SGD + L1)
+    # ------------------------------------------------------------------
+
+    def fit(
+        self,
+        dataset: List[Sample],
+        lr: float = 1e-3,
+        epochs: int = 100,
+        lambda_l1: float = 1e-4,
+        n_steps: int = 50,
+        batch_size: Optional[int] = None,
+        verbose: bool = True,
+    ) -> List[float]:
+        """
+        Stochastic gradient descent with L1 regularisation.
+
+        L1 promotes sparsity — zero parameters mean no coupling between components
+        or inputs, letting the model discover the true dependency structure.
+
+        Returns per-epoch average loss history.
+        """
+        rng = np.random.default_rng(0)
+        loss_history: List[float] = []
+
+        for epoch in range(epochs):
+            indices = rng.permutation(len(dataset))
+            if batch_size is not None:
+                batches: List[np.ndarray] = [
+                    indices[i : i + batch_size]
+                    for i in range(0, len(indices), batch_size)
+                ]
+            else:
+                batches = [indices]
+
+            epoch_loss = 0.0
+            for batch_idx in batches:
+                results = [self.compute_gradient(dataset[i], n_steps) for i in batch_idx]
+                grads, losses = zip(*results)
+                grad = np.mean(grads, axis=0)
+                batch_loss = float(np.mean(losses))
+
+                # L1 subgradient
+                params = self._get_params()
+                grad = grad + lambda_l1 * np.sign(params)
+
+                # Gradient clipping to prevent exploding updates
+                grad_norm = float(np.linalg.norm(grad))
+                if grad_norm > 1.0:
+                    grad = grad / grad_norm
+
+                self._set_params(params - lr * grad)
+
+                # Re-apply mask after update: forbidden positions must stay at 0
+                # even if numerical noise crept in through the L1 subgradient.
+                self.theta_I  *= self.A
+                self.theta_II *= self.A
+
+                epoch_loss += batch_loss
+
+            epoch_loss /= len(batches)
+            loss_history.append(epoch_loss)
+
+            if verbose and (epoch % max(1, epochs // 10) == 0 or epoch == epochs - 1):
+                print(f"Epoch {epoch:4d}/{epochs}: loss = {epoch_loss:.6f}")
+
+        return loss_history
+
+
+COMPONENT_NAMES: List[str] = [
+    "recoater_blade",
+    "nozzle_plate",
+    "heating_elements",
+    "temperature_sensors",
+    "insulation_panels",
+    "firing_resistors",
+    "cleaning_interface",
+    "recoater_motor",
+    "linear_rail",
+]
+
+INPUT_NAMES: List[str] = [
+    "ambient_temperature_c",
+    "build_chamber_temp_c",
+    "ambient_humidity_pct",
+    "powder_contamination_level",
+    "print_hours",
+    "build_volume_cm3",
+    "recoating_speed_mm_s",
+    "recoating_cycles",
+    "maintenance_level",
+]
+
+if __name__ == "__main__":
+    N, C = 9, 9
+    lambda_rates = 1e-5 * np.ones(9)
+
+    A = np.array([
+        [1, 1, 1, 1, 0, 1, 0, 0, 0],
+        [1, 1, 1, 1, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 0, 0, 0],
+        [0, 0, 0, 1, 1, 1, 1, 0, 0],
+        [1, 0, 0, 1, 1, 1, 0, 0, 0],
+        [0, 0, 0, 0, 1, 0, 1, 1, 1],
+        [0, 0, 0, 0, 0, 0, 1, 1, 1],
+        [0, 0, 0, 0, 0, 0, 1, 1, 1],
+        ])   # fully connected; use build_adjacency_matrix() to enforce graph structure
+    print("Adjacency mask A:")
+    print(A)
+
+    model = DegradationModel(N=N, C=C, lambda_rates=lambda_rates,
+                             adjacency_mask=A, seed=42)
+    print(f"\ntheta_I (masked, forbidden positions = 0):")
+    print(model.theta_I.round(5))
+
+    rng = np.random.default_rng(7)
+    dataset: List[Sample] = Sample.from_json("samples.json")
+
+    s = dataset[0]
+    det = model.simulate(s.X, s.tau, n_steps=50, stochastic=False)
+    sto = model.simulate(s.X, s.tau, n_steps=50, stochastic=True, seed=0)
+
+    print("Deterministic simulation:")
+    print(f"  y_hat = {det[-1]}")
+    print(f"  loss  = {model.loss(det[-1], s.y, s.M):.4f}")
+    print("Stochastic simulation (lambda_rates =", lambda_rates, "):")
+    print(f"  y_hat = {sto[-1]}")
+    print(f"  loss  = {model.loss(sto[-1], s.y, s.M):.4f}")
+
+
+    model.summary(component_names=COMPONENT_NAMES)
+    print("\nFitting:")
+    learning_rates = [1e-5, 2e-5, 1e-4, 2e-4]
+    min_loss = 10
+    best_lr = 0
+    for lr in learning_rates:
+        import time
+        model = DegradationModel(N=N, C=C, lambda_rates=lambda_rates,
+                             adjacency_mask=A, seed=int(time.time()))
+        final_loss = model.fit(dataset, lr=1e-3, epochs=300, lambda_l1=0, n_steps=20, verbose=True)[-1]
+        if final_loss < min_loss:
+            model.save("model.npz")
+            min_loss = final_loss
+            best_lr = lr
+    print(best_lr, min_loss)
+    model.summary(component_names=COMPONENT_NAMES)

predict_from_supabase.py

@@ -0,0 +1,140 @@
+"""
+Given a printer_id and timestamp t, fetch health + operating conditions from
+Supabase and return the PPO-recommended replacement schedule.
+"""
+
+import os
+import numpy as np
+from datetime import datetime
+from supabase import create_client
+from stable_baselines3 import PPO
+from model import DegradationModel
+from scheduling_rl import _ACTION_TABLE, COMPONENT_NAMES
+
+# ---------------------------------------------------------------------------
+# Supabase client
+# ---------------------------------------------------------------------------
+
+_sb = create_client(os.environ["SUPABASE_URL"], os.environ["SUPABASE_KEY"])
+
+# Column order must match the obs vector expected by the PPO (C=7 conditions)
+_CONDITION_COLS = [
+    "ambient_temperature_c",
+    "build_chamber_temp_c",
+    "ambient_humidity_pct",
+    "powder_contamination_level",
+    "build_volume_cm3",
+    "recoating_speed_mm_s",
+    "maintenance_level",
+]
+
+_HEALTH_COLS = [
+    "recoater_blade",
+    "nozzle_plate",
+    "heating_elements",
+    "temperature_sensors",
+    "insulation_panels",
+    "firing_resistors",
+    "cleaning_interface",
+    "recoater_motor",
+    "linear_rail",
+]
+
+# ---------------------------------------------------------------------------
+# Fetch helpers
+# ---------------------------------------------------------------------------
+
+def _fetch_health(printer_id: str, t: datetime) -> np.ndarray:
+    """Snapshot whose time_step_id matches hours elapsed since last_repair."""
+    printer = (
+        _sb.table("printers")
+        .select("last_repair")
+        .eq("id", printer_id)
+        .single()
+        .execute()
+        .data
+    )
+    if not printer:
+        raise ValueError(f"Printer {printer_id} not found")
+
+    last_repair = datetime.fromisoformat(printer["last_repair"])
+    time_step_id = int((t - last_repair).total_seconds() // 3600)
+
+    row = (
+        _sb.table("snapshots")
+        .select(", ".join(_HEALTH_COLS))
+        .eq("id", printer_id)
+        .gte("time_step_id", time_step_id)
+        .order("time_step_id", desc=False)
+        .limit(1)
+        .execute()
+        .data
+    )
+    if not row:
+        raise ValueError(f"No snapshot for printer {printer_id} at time_step_id={time_step_id} (t={t})")
+    return np.array([row[0][c] for c in _HEALTH_COLS], dtype=np.float64)
+
+
+def _fetch_conditions(printer_id: str, t: datetime) -> np.ndarray:
+    """Closest conditions row at or before t."""
+    row = (
+        _sb.table("conditions")
+        .select(", ".join(_CONDITION_COLS))
+        .eq("id", printer_id)
+        .lte("timestamp", t.isoformat())
+        .order("timestamp", desc=True)
+        .limit(1)
+        .execute()
+        .data
+    )
+    if not row:
+        raise ValueError(f"No conditions found for printer {printer_id} at {t}")
+    return np.array([row[0][c] for c in _CONDITION_COLS], dtype=np.float64)
+
+# ---------------------------------------------------------------------------
+# Main prediction
+# ---------------------------------------------------------------------------
+
+def predict_replacements(
+    printer_id: str,
+    t: datetime,
+    *,
+    budget_remaining: float,
+    W: float = 10_000.0,
+    t_hours: float = 0.0,
+    ppo_path: str = "scheduler_ppo",
+    model_path: str = "model.npz",
+) -> dict:
+    DegradationModel.load(model_path)   # validates model exists
+    ppo = PPO.load(ppo_path)
+
+    health = _fetch_health(printer_id, t)
+    X_t    = _fetch_conditions(printer_id, t)
+
+    obs = np.concatenate([health, X_t, [budget_remaining / W], [t_hours]]).astype(np.float32)
+    action, _ = ppo.predict(obs, deterministic=True)
+    bits = _ACTION_TABLE[int(action)]
+
+    to_replace = [COMPONENT_NAMES[i] for i, b in enumerate(bits) if b]
+    return {
+        "printer_id":    printer_id,
+        "timestamp":     t.isoformat(),
+        "health":        dict(zip(COMPONENT_NAMES, health.tolist())),
+        "conditions":    dict(zip(_CONDITION_COLS, X_t.tolist())),
+        "replace":       to_replace,
+        "action_id":     int(action),
+    }
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import json, sys
+    printer_id = sys.argv[1] if len(sys.argv) > 1 else "printer_001"
+    t = datetime.fromisoformat(sys.argv[2]) if len(sys.argv) > 2 else datetime.now()
+    budget = float(sys.argv[3]) if len(sys.argv) > 3 else 10_000.0
+
+    result = predict_replacements(printer_id, t, budget_remaining=budget)
+    print(json.dumps(result, indent=2))

scheduling_rl.py

@@ -0,0 +1,530 @@
+"""
+scheduling_rl.py — PPO-based replacement scheduling for the HP Metal Jet S100.
+
+Environment
+-----------
+    State  : [h(9), X_t(9), budget_remaining/W(1), t_hours(1)]  →  R^20
+    Action : Discrete(512) — one of 2^9 joint component-replacement combinations
+    Reward : +1.0 per survived hour (dt=1); 0.0 on the terminal step
+    Done   : any h_i < HEALTH_THRESHOLD  OR  replacement cost exceeds budget
+
+Algorithm
+---------
+    PPO (Stable-Baselines3) with a 512-output Softmax actor.
+    The Softmax over the full 2^9 action space captures all joint replacement
+    correlations without requiring an autoregressive architecture.
+
+    Actor:  Linear(20→64) → Tanh → Linear(64→64) → Tanh → Linear(64→512)
+    Critic: Linear(20→64) → Tanh → Linear(64→64) → Tanh → Linear(64→1)
+
+    SB3 default orthogonal init + 0.01 final-layer scale → initial policy
+    is near-uniform over all 512 actions (max-entropy start).
+
+Component costs (€)
+--------------------
+    Sourced from analogous industrial parts; HP does not publish official prices.
+    Sources listed in DEFAULT_COSTS docstring below.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch as th
+import gymnasium as gym
+from gymnasium import spaces
+from stable_baselines3 import PPO
+
+from model import DegradationModel
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+COMPONENT_NAMES: list[str] = [
+    "recoater_blade",       # 0
+    "nozzle_plate",         # 1
+    "heating_elements",     # 2
+    "temperature_sensors",  # 3
+    "insulation_panels",    # 4
+    "firing_resistors",     # 5
+    "cleaning_interface",   # 6
+    "recoater_motor",       # 7
+    "linear_rail",          # 8
+]
+
+# Replacement costs in euros per component.
+# Sources:
+#   recoater_blade      — WINK3D EOS M290 ceramic blade ~$160, scaled to HP Metal Jet tier
+#                         https://winking3d.com/product/ceramic-recoater-blade-3/
+#   nozzle_plate        — 6× HP TIJ printheads; HP industrial OEM pricing estimate
+#                         https://3dprintingindustry.com/news/hp-launches-new-metal-jet-s100...
+#   heating_elements    — Industrial resistive heating elements for 180°C chamber
+#                         https://www.sentrotech.com/heating-elements/
+#   temperature_sensors — Industrial PT100/thermocouple bundle (Omega, multiple sensors)
+#                         https://www.omega.co.uk/pptst/T3PROBES.html
+#   insulation_panels   — Ceramic fibre panel section for printer-sized chamber
+#                         https://www.sentrotech.com/ceramic-fiber-insulation/
+#   firing_resistors    — Printhead array embedded resistors, proprietary HP part
+#   cleaning_interface  — Wiper blade + solvent delivery module
+#                         https://digiprint-usa.com/blogs/printhead-guides-tips-digiprint-usa/...
+#   recoater_motor      — 500 W–1 kW industrial servo motor
+#                         https://teknic.com/products/clearpath-brushless-dc-servo-motors/
+#   linear_rail         — THK/HIWIN precision guide rail assembly
+DEFAULT_COSTS = np.array([
+     350.0,   # recoater_blade
+    2500.0,   # nozzle_plate
+    1500.0,   # heating_elements
+     200.0,   # temperature_sensors
+     800.0,   # insulation_panels
+     400.0,   # firing_resistors
+     600.0,   # cleaning_interface
+     900.0,   # recoater_motor
+     600.0,   # linear_rail
+], dtype=np.float64)
+
+HEALTH_THRESHOLD: float = 0.1   # printer fails when any component drops below this
+N_COMPONENTS: int = 9
+N_ACTIONS: int = 2 ** N_COMPONENTS  # 512
+
+# Precompute action → binary replacement vector table once at import time.
+# _ACTION_TABLE[i] is a (9,) float32 array: bit j = 1 means replace component j.
+_ACTION_TABLE: np.ndarray = np.array(
+    [[int(b) for b in format(i, f"0{N_COMPONENTS}b")] for i in range(N_ACTIONS)],
+    dtype=np.float32,
+)
+
+
+# ---------------------------------------------------------------------------
+# Environment
+# ---------------------------------------------------------------------------
+
+class PrinterEnv(gym.Env):
+    """
+    Gymnasium environment wrapping DegradationModel for replacement scheduling.
+
+    Parameters
+    ----------
+    model       : Fitted DegradationModel with N=9 components and C inputs.
+    X_series    : (T_max, C) array of predicted operating-condition vectors,
+                  one row per simulated hour. If the episode outlasts X_series
+                  the last row is repeated.
+    W           : Total budget in euros available for the episode.
+    costs       : (9,) replacement cost per component in euros.
+                  Defaults to DEFAULT_COSTS when None.
+    dt          : Simulation time-step in hours (default 1.0).
+    stochastic  : Include Poisson shock terms from model.lambda_rates (default True).
+    seed        : RNG seed for reproducibility.
+    """
+
+    metadata: dict = {"render_modes": []}
+
+    def __init__(
+        self,
+        model: DegradationModel,
+        X_series: np.ndarray,
+        W: float,
+        costs: np.ndarray | None = None,
+        dt: float = 1.0,
+        stochastic: bool = True,
+        seed: int | None = None,
+    ) -> None:
+        """
+        Initialise the environment and validate inputs.
+
+        Builds the observation and action spaces, stores a reference to the
+        fitted DegradationModel, and sets the episode state to its initial
+        values (all components fully healthy, full budget, t=0).
+
+        Raises ValueError if model.N != 9.
+        """
+        super().__init__()
+
+        if model.N != N_COMPONENTS:
+            raise ValueError(f"DegradationModel must have N={N_COMPONENTS}, got {model.N}")
+
+        self.model = model
+        self.X_series = np.asarray(X_series, dtype=np.float64)  # (T_max, C)
+        self.W = float(W)
+        self.costs = (
+            np.asarray(costs, dtype=np.float64) if costs is not None else DEFAULT_COSTS
+        )
+        self.dt = float(dt)
+        self.stochastic = stochastic
+        self._rng = np.random.default_rng(seed)
+
+        # Observation: h(9) + X_t(C) + budget/W(1) + t_hours(1)
+        obs_dim = N_COMPONENTS + model.C + 2
+        self.observation_space = spaces.Box(
+            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32
+        )
+        self.action_space = spaces.Discrete(N_ACTIONS)
+
+        # Episode state — initialised properly in reset()
+        self._h = np.ones(N_COMPONENTS, dtype=np.float64)
+        self._t: float = 0.0
+        self._budget: float = self.W
+        self._step_idx: int = 0
+
+    # ------------------------------------------------------------------
+    # Core Gymnasium interface
+    # ------------------------------------------------------------------
+
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict | None = None,
+    ) -> tuple[np.ndarray, dict]:
+        """
+        Reset the environment to the start of a new episode.
+
+        All components are restored to full health (h=1), the budget is
+        refilled to W, and the simulation clock is set to t=0. If a seed
+        is supplied the internal RNG is re-seeded for reproducibility.
+
+        Returns
+        -------
+        obs : (obs_dim,) float32 observation vector.
+        info : empty dict (required by the Gymnasium API).
+        """
+        if seed is not None:
+            self._rng = np.random.default_rng(seed)
+
+        self._h = np.ones(N_COMPONENTS, dtype=np.float64)
+        self._t = 0.0
+        self._budget = self.W
+        self._step_idx = 0
+
+        return self._obs(), {}
+
+    def step(
+        self, action: int
+    ) -> tuple[np.ndarray, float, bool, bool, dict]:
+        """
+        Advance the simulation by one hour and apply the chosen replacements.
+
+        The action integer is decoded to a 9-bit binary replacement vector via
+        _ACTION_TABLE. Components whose bit is 1 are replaced (health reset to
+        1.0) before the ODE step is taken, so new components benefit from the
+        full degradation rate of a healthy part.
+
+        Budget is checked before applying replacements. If the total cost of
+        the selected action exceeds the remaining budget the episode terminates
+        immediately with reward 0 and no state change.
+
+        After replacements the health vector is updated with one Euler step:
+
+            h_{t+1} = h_t + f(h_t, X_t) * dt  [- Q * H  if stochastic]
+
+        where Q ~ N(0,1)² and H ~ Poisson(λ_i * dt) are independent per
+        component (matching the stochastic extension in model.py §5).
+
+        Parameters
+        ----------
+        action : int in [0, 511].
+
+        Returns
+        -------
+        obs         : (obs_dim,) float32 observation after the step.
+        reward      : dt (1.0) if the printer survived; 0.0 on failure.
+        terminated  : True when any h_i < HEALTH_THRESHOLD or budget exceeded.
+        truncated   : always False (no fixed time limit).
+        info        : dict with keys t_hours, budget_remaining, min_health
+                      (or termination on budget failure).
+        """
+        bits = _ACTION_TABLE[int(action)]          # (9,) float32 binary vector
+        replacement_cost = float(np.dot(bits, self.costs))
+
+        # Hard budget constraint: action is unaffordable → episode ends
+        if replacement_cost > self._budget:
+            return self._obs(), 0.0, True, False, {
+                "termination": "budget_exceeded",
+                "t_hours": self._t,
+                "budget_remaining": self._budget,
+            }
+
+        # Apply replacements: reset selected components to full health
+        self._budget -= replacement_cost
+        self._h[bits.astype(bool)] = 1.0
+
+        # Euler step: dh = f(h, X_t) * dt  [+ optional stochastic shock]
+        X_t = self._current_X()
+        dh = self.model.f(self._h, X_t) * self.dt
+
+        if self.stochastic:
+            Q = self._rng.standard_normal(N_COMPONENTS) ** 2          # N(0,1)² — event intensity
+            H = self._rng.poisson(self.model.lambda_rates * self.dt)  # Poisson(λ_i · dt)
+            dh -= Q * H
+
+        self._h = self._h + dh
+        self._t += self.dt
+        self._step_idx += 1
+
+        # Failure: any component below threshold
+        failed = bool(np.any(self._h < HEALTH_THRESHOLD))
+        reward = 0.0 if failed else self.dt
+
+        return self._obs(), reward, failed, False, {
+            "t_hours": self._t,
+            "budget_remaining": self._budget,
+            "min_health": float(self._h.min()),
+        }
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _obs(self) -> np.ndarray:
+        """
+        Build the current observation vector.
+
+        Layout (obs_dim = 9 + C + 2):
+            h[0:9]              — component health values (float, may be negative
+                                  after a stochastic shock representing catastrophic
+                                  failure before the done flag is raised)
+            X_t[9:9+C]         — current operating-condition inputs from X_series
+            budget/W [9+C]     — remaining budget as a fraction of the initial
+                                  budget W; always in [0, 1] during a valid episode
+            t_hours  [9+C+1]   — elapsed simulation time in hours
+
+        Returns float32 to match observation_space.dtype.
+        """
+        X_t = self._current_X()
+        return np.concatenate([
+            self._h,                        # 9  — component healths (may be < 0 stochastic)
+            X_t,                            # C  — current operating conditions
+            [self._budget / self.W],        # 1  — remaining budget fraction ∈ [0, 1]
+            [self._t],                      # 1  — elapsed hours
+        ]).astype(np.float32)
+
+    def _current_X(self) -> np.ndarray:
+        """
+        Return the operating-condition vector for the current simulation step.
+
+        Clamps the index to the last row of X_series if the episode outlasts
+        the provided forecast horizon, so the environment never raises an
+        IndexError regardless of episode length.
+        """
+        idx = min(self._step_idx, len(self.X_series) - 1)
+        return self.X_series[idx]
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+def train(
+    env: PrinterEnv,
+    *,
+    total_timesteps: int = 1_000_000,
+    save_path: str = "scheduler_ppo",
+    learning_rate: float = 3e-4,
+    n_steps: int = 2048,
+    batch_size: int = 64,
+    n_epochs: int = 10,
+    gamma: float = 0.99,
+    ent_coef: float = 0.01,
+    verbose: int = 1,
+) -> PPO:
+    """
+    Train a PPO agent on PrinterEnv and save the result to disk.
+
+    Architecture
+    ------------
+    Both the actor and critic share the same MLP topology
+    (two hidden layers of 64 units with Tanh activations) but have
+    separate weights, as is standard in Actor-Critic methods:
+
+        Actor  (π_θ): obs(20) → 64 → 64 → logits(512) → Softmax
+        Critic (V_ϕ): obs(20) → 64 → 64 → scalar(1)
+
+    The 512 Softmax output over the full 2^9 action space implicitly models
+    the joint probability distribution over all replacement combinations,
+    capturing correlations (e.g. replacing component 0 should reduce the
+    probability of also replacing component 1 if they interact positively)
+    without requiring an autoregressive sampling pass.
+
+    Initialisation
+    --------------
+    SB3 applies orthogonal initialisation to all layers (scale √2 for hidden,
+    scale 0.01 for the final policy layer). With scale 0.01 the logits start
+    near zero, so Softmax(~0) ≈ 1/512 — effectively maximum entropy over the
+    action space, ensuring the agent explores broadly before committing.
+
+    Hyperparameters
+    ---------------
+    learning_rate : Adam step size (3e-4 is a reliable PPO default).
+    n_steps       : Rollout length before each PPO update (2048 steps ≈
+                    one to several full episodes depending on episode length).
+    batch_size    : Mini-batch size for gradient updates.
+    n_epochs      : Number of gradient passes over each collected rollout.
+    gamma         : Discount factor; 0.99 weights future rewards heavily,
+                    encouraging the agent to maximise long-term lifespan.
+    ent_coef      : Entropy bonus coefficient; keeps action probabilities from
+                    collapsing to a single action too early in training.
+
+    Parameters
+    ----------
+    env             : Configured PrinterEnv instance.
+    total_timesteps : Total environment steps to train for.
+    save_path       : File path (without .zip) for the saved model.
+    learning_rate   : Adam learning rate.
+    n_steps         : Rollout buffer size (steps per PPO update).
+    batch_size      : SGD mini-batch size.
+    n_epochs        : PPO epochs per rollout.
+    gamma           : Discount factor γ.
+    ent_coef        : Entropy regularisation coefficient.
+    verbose         : SB3 verbosity level (0=silent, 1=info, 2=debug).
+
+    Returns
+    -------
+    PPO : The trained Stable-Baselines3 PPO model, ready for evaluate().
+    """
+    policy_kwargs = dict(
+        net_arch=dict(pi=[64, 64], vf=[64, 64]),
+        activation_fn=th.nn.Tanh,
+    )
+
+    ppo = PPO(
+        policy="MlpPolicy",
+        env=env,
+        learning_rate=learning_rate,
+        n_steps=n_steps,
+        batch_size=batch_size,
+        n_epochs=n_epochs,
+        gamma=gamma,
+        ent_coef=ent_coef,
+        policy_kwargs=policy_kwargs,
+        verbose=verbose,
+    )
+
+    ppo.learn(total_timesteps=total_timesteps)
+    ppo.save(save_path)
+    if verbose:
+        print(f"[train] Model saved to {save_path}.zip")
+    return ppo
+
+
+# ---------------------------------------------------------------------------
+# Evaluation
+# ---------------------------------------------------------------------------
+
+def evaluate(
+    ppo: PPO,
+    env: PrinterEnv,
+    n_episodes: int = 10,
+) -> dict:
+    """
+    Run the trained policy deterministically and return summary statistics.
+
+    Each episode starts from a fresh reset() call. The policy is queried with
+    deterministic=True, meaning the action with the highest probability under
+    the current Softmax distribution is always selected (no sampling noise).
+    This gives a reproducible, greedy estimate of the policy's performance.
+
+    Parameters
+    ----------
+    ppo         : Trained PPO model returned by train().
+    env         : PrinterEnv instance (can be the same env used for training).
+    n_episodes  : Number of evaluation episodes to average over.
+
+    Returns
+    -------
+    dict with keys:
+        mean_hours        — average hours survived across all episodes.
+        std_hours         — standard deviation of hours survived.
+        mean_replacements — average total number of component replacements
+                            performed per episode (summed over all steps).
+        mean_budget_spent — average euros spent on replacements per episode.
+        episodes          — list of per-episode dicts, each containing:
+                            episode, hours_survived, budget_spent, replacements.
+    """
+    hours, n_replacements, budget_spent = [], [], []
+    episodes = []
+
+    for ep in range(n_episodes):
+        obs, _ = env.reset()
+        done = False
+        ep_replacements = 0
+
+        while not done:
+            action, _ = ppo.predict(obs, deterministic=True)
+            obs, _, terminated, truncated, _ = env.step(int(action))
+            done = terminated or truncated
+            ep_replacements += int(_ACTION_TABLE[int(action)].sum())
+
+        hours.append(env._t)
+        n_replacements.append(ep_replacements)
+        budget_spent.append(env.W - env._budget)
+        episodes.append({
+            "episode":        ep,
+            "hours_survived": env._t,
+            "budget_spent":   env.W - env._budget,
+            "replacements":   ep_replacements,
+        })
+
+    return {
+        "mean_hours":        float(np.mean(hours)),
+        "std_hours":         float(np.std(hours)),
+        "mean_replacements": float(np.mean(n_replacements)),
+        "mean_budget_spent": float(np.mean(budget_spent)),
+        "episodes":          episodes,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Sanity demo
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+
+    N, C = 9, 9
+    lambda_rates = np.array([0.05, 0.10, 0.02, 0.03, 0.01, 0.08, 0.06, 0.04, 0.03])
+
+    deg_model = DegradationModel(
+        N=N, C=C, lambda_rates=lambda_rates, seed=0
+    )
+
+    # Synthetic X_series: 8 000 hours of operating conditions drawn from [0, 1]
+    rng = np.random.default_rng(42)
+    #X_series = rng.uniform(0.0, 1.0, size=(8_000, C))
+
+    env = PrinterEnv(
+        model=deg_model,
+        X_series=X_series,
+        W=10_000.0,
+        dt=1.0,
+        stochastic=True,
+        seed=42,
+    )
+
+    print("Observation space:", env.observation_space.shape)
+    print("Action space:     ", env.action_space.n, "discrete actions")
+    print()
+
+    # --- Random-policy baseline ---
+    obs, _ = env.reset(seed=0)
+    done = False
+    while not done:
+        action = env.action_space.sample()
+        obs, _, terminated, truncated, info = env.step(action)
+        done = terminated or truncated
+    print(f"Random policy  → {env._t:.0f} h survived | €{env.W - env._budget:.0f} spent")
+
+    # --- PPO training (short demo: 200k steps) ---
+    print("\nTraining PPO (200 000 timesteps) ...")
+    trained = train(
+        env,
+        total_timesteps=200_000,
+        save_path="scheduler_ppo",
+        verbose=1,
+    )
+
+    # --- Evaluation ---
+    results = evaluate(trained, env, n_episodes=10)
+    print(f"\n{'─'*45}")
+    print(f"  Mean hours survived : {results['mean_hours']:>8.1f} h")
+    print(f"  Std  hours          : {results['std_hours']:>8.1f} h")
+    print(f"  Mean replacements   : {results['mean_replacements']:>8.1f}")
+    print(f"  Mean budget spent   : €{results['mean_budget_spent']:>7.0f}")
+    print(f"{'─'*45}")