Skip to main content
By default, Onyx models predict a single point estimate per output, trained with mean squared error (MSE) loss. With uncertainty estimation enabled, the model instead predicts both a mean and a variance per output, giving you a measure of confidence in each prediction.

​
How It Works

​
Standard Training (MSE Loss)

Without uncertainty estimation, the model minimizes the mean squared error between its predictions and the true values: LMSE=1N∑i=1N(yi−μi)2L_{\text{MSE}} = \frac{1}{N} \sum_{i=1}^{N} (y_i - \mu_i)^2 where μi\mu_i is the model’s prediction and yiy_i is the true value.

​
Uncertainty Training (Gaussian NLL Loss)

With predict_uncertainty=True, the model outputs both a mean μ\mu and a log-variance log⁡(σ2)\log(\sigma^2) for each output. Training uses the Gaussian negative log-likelihood loss: LNLL=1N∑i=1N[log⁡(σi2)+(yi−μi)2σi2]L_{\text{NLL}} = \frac{1}{N} \sum_{i=1}^{N} \left[ \log(\sigma_i^2) + \frac{(y_i - \mu_i)^2}{\sigma_i^2} \right] Intuition: The model is penalized for both prediction error and miscalibrated confidence. If the model is uncertain about a prediction, it can increase σ2\sigma^2 to reduce the squared-error term — but it pays a penalty via the log⁡(σ2)\log(\sigma^2) term. This forces the model to be honest about what it knows vs. doesn’t know.

​
Enabling Uncertainty Estimation

Set predict_uncertainty=True on any model configuration. This works with all model types (MLP, RNN, Transformer). 
from onyxengine.modeling import MLPConfig, Input, Output

config = MLPConfig(
    outputs=[Output(name='acceleration')],
    inputs=[
        Input(name='velocity', parent='acceleration', relation='derivative'),
        Input(name='position', parent='velocity', relation='derivative'),
        Input(name='control_input'),
    ],
    dt=0.01,
    sequence_length=8,
    hidden_layers=3,
    hidden_size=64,
    predict_uncertainty=True
)
Training is automatic — the backend switches to Gaussian NLL loss when predict_uncertainty=True.

​
Getting Uncertainty from Direct Inference

When uncertainty is enabled, model(x) returns a tuple (mean, variance) instead of a single tensor. Both are in real-world units — no manual unscaling is needed: 
import torch
from onyxengine import Onyx

onyx = Onyx()
model = onyx.load_model('my_uncertainty_model')

# Create input tensor
batch_size = 1
seq_len = model.config.sequence_length
num_inputs = len(model.config.inputs)
x = torch.randn(batch_size, seq_len, num_inputs)

with torch.no_grad():
    mean, variance = model(x)

print("Mean prediction:", mean)       # Shape: (batch_size, num_direct_outputs)
print("Variance:", variance)           # Shape: (batch_size, num_direct_outputs)

# Get standard deviation
std = torch.sqrt(variance)

print("Predicted mean:", mean)
print("Predicted std:", std)

​
Getting Uncertainty from simulate()

When simulating with an uncertainty-enabled model, the SimulationResult includes an output_uncertainties attribute containing the variance at each time step, already unscaled to real-world units: 
import torch
import numpy as np
import matplotlib.pyplot as plt
from onyxengine import Onyx

onyx = Onyx()
model = onyx.load_model('my_uncertainty_model')

# Set up simulation
batch_size = 1
seq_len = model.config.sequence_length
sim_steps = 200

velocity = torch.zeros(batch_size, seq_len, 1)
position = torch.zeros(batch_size, seq_len, 1)
control_input = torch.ones(batch_size, seq_len + sim_steps, 1)

# Run simulation
result = model.simulate(
    x0={'velocity': velocity, 'position': position},
    external_inputs={'control_input': control_input},
    sim_steps=sim_steps
)

# output_uncertainties is a FeatureTrajectory (or None if model wasn't trained with uncertainty)
# Values are variance in real-world units
accel_variance = result.output_uncertainties['acceleration']  # Shape: (batch_size, sim_steps, 1)
accel_std = torch.sqrt(accel_variance)

# Plot mean prediction with confidence bands
time = np.arange(sim_steps) * model.config.dt
accel_mean = result.outputs['acceleration'][0, :, 0].cpu().numpy()
accel_std_np = accel_std[0, :, 0].cpu().numpy()

plt.figure(figsize=(10, 4))
plt.plot(time, accel_mean, label='Mean prediction')
plt.fill_between(
    time,
    accel_mean - 2 * accel_std_np,
    accel_mean + 2 * accel_std_np,
    alpha=0.3,
    label='±2σ confidence'
)
plt.xlabel('Time (s)')
plt.ylabel('Acceleration')
plt.legend()
plt.title('Simulation with Uncertainty Bands')
plt.show()

​
Key Notes

  • Only direct outputs have uncertainty estimates. Derived outputs (those computed via parent/relation, like velocity integrated from acceleration) are deterministic functions of their parents and do not have separate uncertainty predictions.
  • Simulation uses means for state integration. During simulate(), only the mean predictions are used to update states. Variance is tracked separately and does not affect the trajectory rollout.
  • This captures aleatoric uncertainty. The learned variance reflects inherent noise in the data (aleatoric uncertainty), not uncertainty due to limited training data or model capacity (epistemic uncertainty). Regions with noisier data will naturally have higher predicted variance.