Training a Battery AI Model with PyBaMM: Predicting SOH and RUL

The previous article built the full pipeline for generating physics-labelled EIS data with PyBaMM. This one moves to the next step: training sequence models (LSTM, Transformer) to predict state of health (SOH) and remaining useful life (RUL).

One thing has to be said first. SOH and RUL are tasks of completely different difficulty, and treating them as one multi-task problem is the most common form of self-deception in this field. SOH estimates the present state, and the information for it is already in the input. RUL extrapolates the future, and early in a cell’s life it is fundamentally ill-posed. The next sections explain why, and what to do about it.

PyBaMM battery modelling series (4 parts): ① Architecture and solvers → ② EIS labels → ③ Dataset pipeline → ④ Training SOH and RUL. This is part 4.

1. Degradation as a stochastic process

Battery ageing is non-Markovian: internal state at cycle $k$ depends on the entire stress history. Under optimal estimation, SOH and RUL are hidden states of a discrete-time nonlinear system. The classical formulation is an extended Kalman filter:

State transition (SEI growth and loss of active material):

$$ x_{k+1} = f(x_k, u_k) + w_k, \quad w_k \sim \mathcal{N}(0, Q) $$

Observation (EIS spectra and terminal voltage):

$$ y_k = h(x_k, u_k) + v_k, \quad v_k \sim \mathcal{N}(0, R) $$

$x_k$ holds the internal capacity parameters that map directly to SOH (LLI, LAM). The deep learning task is to replace the heuristic observation function $h(\cdot)$ with a differentiable network $\theta$ that maximises the joint log-likelihood of the degradation trajectory.

There is an under-appreciated corollary here: a network fed only the current cycle’s EIS is structurally incapable of expressing non-Markovian behaviour. It needs a history window, or no amount of depth turns it into more than a noisy lookup table.

2. Why RUL is ill-posed early in life

Take two cells whose capacity curves are nearly identical for the first 200 cycles. One hits accelerated lithium plating — a knee point — at cycle 400; the other fades smoothly to 800. At cycle 100 their observables are almost indistinguishable while their true RUL differs by a factor of two.

This is not a modelling capacity problem. The information is simply not present in the input. Force the network to emit a point estimate anyway and it regresses to the conditional mean of the training set — which shows up early in life as predictions hugging the dataset’s average lifetime. The metric looks acceptable; the discriminative power is nil.

Three responses, in order of preference:

What to avoid is plain MSE point regression on RUL reported as a single average error — that number hides complete early-life failure.

3. A model definition that actually runs

The following is self-contained. Positional encoding has to be written out — PyTorch ships no PositionalEncoding class, and a great many tutorials reference one that does not exist.

import math
import torch
import torch.nn as nn

class PositionalEncoding(nn.Module):
    """Standard sinusoidal encoding. "Position" here is cycle index, not
    wall-clock time - which only works for evenly spaced cycles. If your EIS
    is measured at irregular intervals, feed delta-cycle as a feature instead
    of relying on positional encoding."""

    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))   # buffer, not parameter: no
                                                      # gradient, but saved with the model

    def forward(self, x):
        return x + self.pe[:, : x.size(1)]


class ImpedanceTransformer(nn.Module):
    def __init__(self, input_dim=120, d_model=256, nhead=8, num_layers=4, n_quantiles=3):
        super().__init__()
        self.input_projection = nn.Linear(input_dim, d_model)
        self.pos_encoder = PositionalEncoding(d_model)

        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=4 * d_model,
            batch_first=True,        # without this the convention is [seq, batch, dim]
                                     # and silently misaligns with the [batch, seq, dim] below
            norm_first=True,         # pre-LN, far more stable when deep
        )
        self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers)

        self.soh_head = nn.Linear(d_model, 1)
        # RUL emits quantiles rather than a point estimate - see the previous section
        self.rul_head = nn.Sequential(
            nn.Linear(d_model, 64), nn.GELU(), nn.Linear(64, n_quantiles), nn.Softplus()
        )

    def forward(self, eis_sequence, pad_mask=None):
        # eis_sequence: [batch, seq_len, input_dim]
        # pad_mask:     [batch, seq_len], True marks padding
        x = self.pos_encoder(self.input_projection(eis_sequence))
        feats = self.transformer(x, src_key_padding_mask=pad_mask)

        if pad_mask is None:
            latest = feats[:, -1, :]
        else:
            # For variable-length sequences [-1] is padding. Take the last valid step.
            lengths = (~pad_mask).sum(dim=1) - 1
            latest = feats[torch.arange(feats.size(0)), lengths]

        # Quantiles must be monotone: cumulative sum guarantees q10 <= q50 <= q90
        rul = torch.cumsum(self.rul_head(latest), dim=-1)
        return self.soh_head(latest), rul

Three silent failure modes are flagged in the comments: omitting batch_first misaligns the dimension convention without raising anything; taking [:, -1, :] on padded sequences reads padding; and unconstrained quantile heads happily produce q90 below q10.

4. Loss: put the physics in, do not check it afterwards

$$ \mathcal{L}(\theta) = \lambda_1 \| \text{SOH}_{pred} - \text{SOH}_{true} \|_2^2 + \lambda_2 \mathcal{L}_{pinball}(\text{RUL}) + \lambda_3 \Phi(x) $$

import torch
import torch.nn.functional as F

QUANTILES = torch.tensor([0.1, 0.5, 0.9])

def pinball_loss(pred, target, quantiles=QUANTILES):
    """pred: [batch, n_q]  target: [batch, 1]"""
    q = quantiles.to(pred.device).view(1, -1)
    err = target - pred
    return torch.maximum(q * err, (q - 1) * err).mean()

def monotonicity_penalty(soh_seq):
    """SOH must be non-increasing over cycles. Rest recovery causes small genuine
    rebounds, so allow a tolerance band and only penalise rises beyond it - a hard
    constraint would suppress real physics."""
    diff = soh_seq[:, 1:] - soh_seq[:, :-1]
    return F.relu(diff - 0.002).pow(2).mean()

def total_loss(soh_pred, rul_pred, soh_true, rul_true, soh_seq,
               l1=1.0, l2=0.5, l3=0.1):
    return (l1 * F.mse_loss(soh_pred, soh_true)
            + l2 * pinball_loss(rul_pred, rul_true)
            + l3 * monotonicity_penalty(soh_seq))

The penalty uses relu(diff - 0.002) rather than relu(diff) because post-rest capacity recovery is real. Treating it as a violation teaches the model an over-smoothed curve that goes blunt exactly where the knee is.

5. Covariate shift and leakage

The deadliest mistake with synthetic data is a row-wise random split. Time series derived from a single Simulation.solve() trajectory have fully deterministic covariance; if cycle $N$ and $N+1$ of the same cell land on opposite sides of the split, the Transformer can cheat by interpolation and then fail catastrophically on real hardware.

import numpy as np
from sklearn.model_selection import GroupKFold

groups = metadata["cell_design_id"].values
gkf = GroupKFold(n_splits=5)

for fold, (tr, te) in enumerate(gkf.split(X, y, groups=groups)):
    # Assert in the pipeline. Discovering implausibly good metrics later costs more.
    assert not (set(groups[tr]) & set(groups[te])), f"group leak in fold {fold}"
    ...

There is a subtler leak: feature normalisation. Computing mean and standard deviation over the full dataset before splitting bleeds test statistics into training. Normalisation parameters must come from the training fold alone, then be applied to the test fold.

6. Evaluation: one average error will lie to you

What makes battery life prediction hard is early-life uncertainty, error near the knee, and cross-cell generalisation. Reporting only RUL MAE conceals all three.

Quantile coverage deserves particular attention. A model with excellent RUL MAE whose [q10, q90] interval covers only 45% of truths is more dangerous in production than one with twice the error and honest intervals — it invites you to trust it at exactly the wrong moment.

7. Deploying to a BMS: quantisation is not free

A real BMS runs on an ASIL-rated automotive MCU below 100 MHz. A multi-head attention model trained in PyTorch has to be compressed hard: post-training quantisation (PTQ) to int8 plus structured pruning, exported through ONNX, running alongside the hardware EKF matrix unit.

But int8 PTQ carries a trap that must be validated for. We measured this on a separate project: an int8 model that was essentially lossless on the standard validation set collapsed on out-of-distribution samples. Validating quantisation only on well-behaved data badly overstates how safe it is.

Mapped onto batteries, quantisation validation has to cover these regions separately:

• Samples near the knee — the steepest gradients, where quantisation error is amplified most
• Low-temperature operation — spectra differ markedly from room temperature and activation distributions shift
• Degradation pathways rare in the training set, such as plating-dominated rather than SEI-dominated cells

The calibration set must include these regions; otherwise PTQ fits its dynamic range to the common cases and the rare ones saturate outright. If int8 proves unacceptable there, fall back to fp16 — half the size benefit for predictable behaviour is usually the right trade.

8. Experiment record template

Data source:    PyBaMM DFN/SPMe + measured calibration set
Parameter set:  OKane2022 (degradation) / self-fitted
Split:          GroupKFold by cell_design_id / protocol_id
Normalisation:  computed on training fold only, applied to test fold
Input window:   last N EIS spectra + temperature + DOD + C-rate
Targets:        SOH (point), RUL (q10/q50/q90)
Key metrics:    SOH MAE, stage-wise RUL MAE, quantile coverage, knee point error
Leak check:     assert no group overlap between train and test
Quant check:    knee / low-temp / rare pathway reported separately
Deployment:     ONNX / int8 or fp16 / MCU latency budget

This record is what connects physical modelling, machine learning evaluation and embedded constraints. Without it, a model scoring well on a random split has most likely memorised synthetic trajectories rather than learned transferable ageing behaviour.

References

• Python Battery Mathematical Modelling (PyBaMM)
• PyTorch TransformerEncoder
• PyBOP: battery model optimisation and parameterisation
• Synthetic dataset of LG M50 batteries with different degradation pathways