How do I smooth out my volatile telemetry predictions?

[AyushG-1210]

Problem Context:
I am training a Hybrid MoE-PINN to discover the local grip limit ($\mu$) of a race track based on F1 telemetry. The model uses two experts (Straight vs. Corner) gated by curvature. Despite using L-BFGS and Soft-Gating, the model is stuck in a local minimum where the grip $(\mu)$ refuses to "self-discover" beyond the hard-coded floor, and the resulting telemetry (throttle/brake) exhibits high-frequency oscillations.

Code:
This is a pretty condensed version of my actual scripts, but the main parts should be clear enough to diagnose the error.

import torch
import torch.nn as nn
import torch.nn.functional as F
import deepxde as dde
import numpy as np

class StraightExpert(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(6, 64), nn.Tanh(),
            nn.Linear(64, 32), nn.Tanh()
        )
        self.out = nn.Linear(32, 1)

    def forward(self, x):
        return 1.2 + F.softplus(self.out(self.net(x))) * 0.3

class CornerExpert(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(7, 64), nn.Tanh(),
            nn.Linear(64, 32), nn.Tanh()
        )
        self.out = nn.Linear(32, 1)

    def forward(self, x):
        return 1.2 + F.softplus(self.out(self.net(x))) * 0.3

class MoEPINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.straight_net = StraightExpert()
        self.corner_net = CornerExpert()
        self.max_power_watts = 750000.0
        self.max_traction_newtons = 22000.0
        self.brake_force_newtons = 45000.0

    def get_engine_force(self, velocity_ms):
        power_limit = self.max_power_watts / (velocity_ms + 1.0)
        return torch.min(power_limit, torch.tensor(self.max_traction_newtons, device=velocity_ms.device))

    def gate_function(self, curvature):
        return torch.sigmoid((torch.abs(curvature) - 0.002) * 100)

    def forward(self, x):
        x_phys = x[:, 0:6]
        curvature = x[:, 16:17]
        mu_s = self.straight_net(x_phys)
        mu_c = self.corner_net(torch.cat([x_phys, curvature], dim=1))
        w = self.gate_function(curvature)
        return (1 - w) * mu_s + w * mu_c

def get_physics_pde(moepinn_model):
    def pde(x, y):
        mass = torch.tensor(798.0, device=y.device)
        g_force = mass * 9.81
        curvature, throttle, brake, v_ms = x[:, 16:17], x[:, 17:18], x[:, 18:19], x[:, 19:20]
        
        v_anchor = torch.sqrt((1.4 * 9.81) / torch.clamp(torch.abs(curvature), min=1e-5))
        res_v_anchor = torch.relu(v_anchor - v_ms) * 0.1
        res_grip_reward = ((y - 1.45) ** 2) * 0.05
        
        f_eng = moepinn_model.get_engine_force(v_ms) * throttle
        f_drag = 0.5 * 1.225 * 1.6 * (v_ms ** 2)
        res_flux = torch.relu(f_drag - (0.90 * f_eng))
        
        f_net = f_eng - (45000.0 * brake) - f_drag - (0.015 * g_force)
        f_lat = mass * (v_ms ** 2) * torch.abs(curvature)
        f_long = (mass * (f_net / mass)) + f_drag + (0.015 * g_force)
        
        target_g = (torch.sqrt(f_lat**2 + f_long**2) * 1.05) / g_force
        supply_g = (y * (g_force + 0.5 * 1.225 * 3.0 * 1.6 * v_ms**2)) / g_force
        res_fit = (supply_g - target_g) ** 2

        dmu_ds = dde.grad.jacobian(y, x, j=16)
        d2mu_ds2 = dde.grad.hessian(y, x, i=16, j=16)

        mu_prev_s = moepinn_model.straight_net(x[:, 6:12])
        mu_prev_c = moepinn_model.corner_net(torch.cat([x[:, 6:12], x[:, 7:8]], dim=1))
        mu_prev = (1 - moepinn_model.gate_function(x[:, 7:8])) * mu_prev_s + moepinn_model.gate_function(x[:, 7:8]) * mu_prev_c
        res_hist = ((y - mu_prev) ** 2) * 0.5

        return [res_fit, torch.relu(0.15 - throttle), torch.relu(brake - 0.95), dmu_ds * 0.1, d2mu_ds2 * 0.1, res_v_anchor, res_grip_reward, res_flux, res_hist]
    return pde

def main():
    X_train = np.random.rand(1000, 20).astype(np.float32)
    geom = dde.geometry.Hypercube(xmin=[0]*20, xmax=[1]*20)
    net = MoEPINN()
    data = dde.data.PDE(geometry=geom, pde=get_physics_pde(net), bcs=[], num_domain=0, anchors=X_train)
    model = dde.Model(data, net)
    
    weights = [15.0, 12.0, 12.0, 8.0, 12.0, 4.0, 3.0, 5.0, 10.0]
    model.compile("adam", lr=0.001, loss_weights=weights)
    model.train(iterations=1000)
    
    dde.optimizers.config.set_LBFGS_options(maxiter=500)
    model.compile("L-BFGS", loss_weights=weights)
    model.train()

if __name__ == "__main__":
    main()

Nothing I seem to do fixes the car trying to jump from max throttle to max braking.

[echen5503]

1. You mention that the model can't discover past a floor. This behavior is mathematically forced by your code. Note that in return 1.2 + F.softplus(self.out(self.net(x))) * 0.3, soft plus is always > 0.

2. Note that multiplying by 100 will make an extremely steep gradient, that could negatively affect performance. You could reduce it to make it more gradual swapping from max throttle to max braking

def gate_function(self, curvature):
    return torch.sigmoid((torch.abs(curvature) - 0.002) * 100)

3. Note that there is a caching bug with torch IC/BC caching causes critical errors in inverse problems for IC/BC #2074, which likely affects your code. Check the PR to see how to fix this.

[echen5503]

(I assume that you left out some code for boundary conditions in your code.)