Simulating Models

Onyx models come with built-in simulation capabilites to quickly deploy trained models.

Model Simulator

To recap from Training Models, AI models predict one step at a time by default. But for simulation and controls, we need to predict more than a single step and instead simulate trajectories over multiple steps. Managing state relations, numerical integration, and recursive model calls is tedious, so we provide a ModelSimulator to handle this for you efficiently.

All Onyx models inherit from a ModelSimulator class that adds a simulate() method to the model. To use this method, we just need to configure the model’s simulator.

Once the model is trained/saved, the simulator config will be saved with it, so you will not need to set the simulator when loading the model.

from onyxengine.modeling import (
    ModelSimulatorConfig,
    State
)
import onyxengine as onyx

# Configure a model simulator for a new model
sim_config = ModelSimulatorConfig(
    outputs=['acceleration'],
    states=[
        State(name='velocity', relation='derivative', parent='acceleration'),
        State(name='position', relation='derivative', parent='velocity'),
    ],
    controls=['control_input'],
    dt=0.0025
)

Sim configs are similar to dynamical system state-space models, where \(\dot{x} = f(x, u)\). In this example:

Model Outputs
- acceleration (\(\dot{x}\)) is the output of the model.
Model Inputs
- velocity (\(x_1\)) is the first state, whose parent acceleration is its derivative.
- position (\(x_2\)) is the second state, whose parent velocity is its derivative.
- control_input (\(u\)) is a control input and provided at each time step.

The control inputs are given at each time step, and outputs are predicted by the model, so they do not need any special attention.

The states, however, must be calculated from the outputs (or other states) and fed back into the model as inputs to rollout a trajectory. The available relations between states and their parents are:

“output”: The state is equal to an output of the model.
- \(\text{state}_{t+1} = \text{parent}_t\)
“delta”: The parent is the delta of the state.
- \(\text{state}_{t+1} = \text{state}_t + \text{parent}_t\)
“derivative”: The parent is the derivative of the state.
- \(\text{state}_{t+1} = \text{state}_t + \text{parent}_t \cdot \text{dt}\)

Note

Often it can be easier for a hardware AI model to predict the delta or derivative of a feature than the feature itself.

Lastly, the dt parameter is the model’s time step, this parameter is used both in simulation and training.

simulate()

We will now use the simulate() method on a model. Here is the example code:

import torch
import onyxengine as onyx

# Load our model
model = onyx.load_model('example_model')
num_states = model.config.sim_config.num_states
num_controls = model.config.sim_config.num_controls
num_inputs = model.config.sim_config.num_inputs
seq_length = model.config.sequence_length

# Run inference with our model (using normal pytorch model prediction)
batch_size = 1
test_input = torch.ones(batch_size, seq_length, num_inputs)
with torch.no_grad():
    test_output = model(test_input)
print(test_output)

# Simulate a trajectory with our model
# Model will fill in the traj_solution tensor with the simulated trajectory
sim_steps = 10
x0 = torch.ones(batch_size, seq_length, num_states)
u = torch.ones(batch_size, sim_steps, num_controls)
traj_solution = torch.zeros(1, sim_steps, num_inputs)
model.simulate(traj_solution, x0, u)
print(traj_solution)

To make the simulate() method performant, we avoid using dynamic memory allocation by passing in the pre-allocated tensors needed:

traj_solution - The solution tensor whose states will be “filled in” with the simulated trajectory.
- Should be of shape (batch_size, sim_steps, num_inputs):
  batch_size is the number of parallel trajectories to simulate
  
  sim_steps is the number of time steps to simulate
  
  num_inputs is the number of inputs to the model (num_states + num_controls).
x0 - The initial state for the simulation, and the first sequence_length state values of the traj_solution tensor.
- Should be of shape (batch_size, seq_length, num_states)
u - The control inputs for the simulation, and the complete array of control values for the traj_solution.
- Should be of shape (batch_size, sim_steps, num_controls)
(Optional) output_traj - Optionally, you can pass a tensor to store the model outputs
- Should be of shape (batch_size, sim_steps, num_outputs)

The batch dimension allows for parallel simulation of multiple trajectories, which is where GPU acceleration becomes useful.

If you need to run some code (such as a controller) at each time step, you can just use the simulate method in a loop and simulate one step at a time.