Skip to main content
This guide walks you through training a model using the example dataset included with every new Onyx Engine account.
Make sure you’ve completed the installation steps before continuing.

​
Train a Model

Create a new Python file and add the following code: 
from onyxengine import Onyx
from onyxengine.modeling import (
    Output,
    Input,
    MLPConfig,
    TrainingConfig,
    AdamWConfig,
)

# Initialize the client (defaults to ONYX_API_KEY env var)
onyx = Onyx()

# Define model outputs and inputs
outputs = [
    Output(name='acceleration'),
]
inputs = [
    Input(name='velocity', parent='acceleration', relation='derivative'),
    Input(name='position', parent='velocity', relation='derivative'),
    Input(name='control_input'),
]

# Configure the model
model_config = MLPConfig(
    outputs=outputs,
    inputs=inputs,
    dt=0.0025,
    sequence_length=8,
    hidden_layers=3,
    hidden_size=64,
    activation='relu',
    dropout=0.2,
    bias=True
)

# Configure training
training_config = TrainingConfig(
    training_iters=3000,
    train_batch_size=1024,
    test_dataset_size=500,
    checkpoint_type='multi_step',
    optimizer=AdamWConfig(lr=3e-4, weight_decay=1e-2),
    lr_scheduler=None
)

# Start training
onyx.train_model(
    model_name='my_first_model',
    model_config=model_config,
    dataset_name='example_data',
    training_config=training_config,
)
Run the script to start training. You can view the trained model in the platform. 
Preparing to train model [my_first_model] using dataset [example_data].

​
Use the Trained Model

Once training completes, you can pull the model from code and simulate a trajectory: 
import torch
from onyxengine import Onyx

# Initialize the client
onyx = Onyx()

# Load the trained model
model = onyx.load_model('my_first_model')
batch_size = 1
sequence_length = model.config.sequence_length
sim_steps = 10

# Initial conditions (velocity and position)
velocity = torch.zeros(batch_size, sequence_length, 1)
position = torch.zeros(batch_size, sequence_length, 1)

# External inputs for the full simulation trajectory (control input)
control_input = torch.ones(batch_size, sequence_length + sim_steps, 1)

# Simulate the trajectory, each timestep will be model.config.dt seconds
result = model.simulate(
    x0={"velocity": velocity, "position": position},
    external_inputs={"control_input": control_input},
    sim_steps=sim_steps,
)

# Access results
print("Velocity trajectory shape:", result.inputs['velocity'].shape)
print("Position trajectory shape:", result.inputs['position'].shape)
print("Acceleration output shape:", result.outputs['acceleration'].shape)

# View the model outputs
print("Acceleration output:", result.outputs['acceleration'])

​
Understanding the Code

​
Model Inputs and Outputs

Onyx comes with a set of built-in tools to help create and manage model features:
  • Output: Features the model predicts (acceleration)
  • Input: Features fed into the model (velocity, position, control_input)
The parent and relation parameters of the Input and Output classes define how features update in the simulate() method:
  • 'derivative': The parent is the time derivative of this feature
  • 'delta': The parent is the delta change per timestep
  • 'equal': The parent is equal to this feature (feed output straight back in as input)
In this example, we chained acceleration -> velocity -> position. You can also always use the underlying PyTorch model as is without calling simulate(): 
# Call the PyTorch model prediction directly
test_input = torch.ones(batch_size, sequence_length, num_features)
with torch.no_grad():
    test_output = model(test_input)
print(test_output)

​
Simulation

The simulate() method handles multi-step prediction:
  1. Takes initial state values in x0
  2. Takes the full trajectory for non-derived features in external_inputs
  3. Automatically integrates states using the defined relations to roll out the trajectory
  4. Returns a SimulationResult with trajectory data for all features
Note that simulating with integration schemes is a serial operation, since each timestep depends on the previous one. For applications that can leverage parallel GPU acceleration, you may want to avoid features that require step-wise integration (no parent and relation).

​
Next Steps

Uploading Datasets

Collect, save, and load datasets in the Onyx Engine

Training Models

Deep dive into model configuration and training options

Optimizing Models

Automatically search for the best model architecture and hyperparameters

Simulating with Models

Learn advanced simulation techniques