Model Predictive Control (MPC)¶

MPC uses a dynamics model to predict system behavior and choose an optimal control sequence under constraints. At each step it solves an optimization problem, applies the first control input from the optimal sequence, and repeats with a shifted horizon.

Theory (brief)¶

Discrete dynamics: \(x_{k+1} = f(x_k, u_k)\)
Horizon cost \(N\):

\[ J = \sum_{i=0}^{N-1} (x_{k+i} - x^{\mathrm{ref}}_{k+i})^\top Q (x_{k+i} - x^{\mathrm{ref}}_{k+i}) + u_{k+i}^\top R\, u_{k+i} + \Delta u_{k+i}^\top S\, \Delta u_{k+i} + \text{terminal\_weight} \cdot (x_{k+N}-x^{\mathrm{ref}}_{k+N})^\top Q (x_{k+N}-x^{\mathrm{ref}}_{k+N}) \]

Control increment:

\[ \Delta u_{k+i} = u_{k+i} - u_{k+i-1} \]

Constraints:

\[ \begin{aligned} u_{\min} \le u_{k+i} \le u_{\max}, &\quad \Delta u_{\min} \le \Delta u_{k+i} \le \Delta u_{\max}, \\ \end{aligned} \]

Receding horizon: solve → apply \(u_k\) → shift window → repeat
Stability: terminal weight, sufficient \(N\), feasibility

Architecture¶

The MPC module consists of:

Component	Class	Description
Low-level solver	`MPC`	Projected-gradient optimizer over a differentiable dynamics
High-level agent	`MPCAgent`	DSAC-like wrapper with learned dynamics, buffer, training
Weights config	`MPCWeights`	Diagonal Q, R, S weights and terminal weight
Constraints	`MPCConstraints`	Box constraints for u and du
Extra costs	`MPCTrackingExtraCostConfig`, `MPCStepResponseExtraCostConfig`	Additional penalties for smoothness, overshoot, settling
Dynamics models	`OneStepMLP`, `NARXDynamicsModel`, `TransformerDynamicsModel`	Neural network models for learning dynamics
Scaler	`MPCStandardScaler`	Feature normalization (mean/std)

Quick Start¶

Basic MPC with a custom dynamics function¶

import numpy as np
import torch
from tensoraerospace.agent.mpc import MPC, MPCWeights, MPCConstraints

state_dim = 4
action_dim = 1

# Define dynamics: x_{t+1} = f(x_t, u_t)
def dynamics(x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
    # Simple linear dynamics for example
    A = torch.eye(state_dim)
    B = torch.zeros(state_dim, action_dim)
    B[-1, 0] = 1.0  # control affects last state
    return x @ A.T + u @ B.T

# Configure weights
weights = MPCWeights(
    Q_diag=np.array([1.0, 1.0, 1.0, 10.0]),  # state tracking weights
    R_diag=np.array([0.01]),                   # control effort
    S_diag=np.array([0.1]),                    # control smoothness
    terminal_weight=2.0,
)

# Configure constraints
constraints = MPCConstraints(
    u_min=np.array([-1.0]),
    u_max=np.array([1.0]),
    du_min=np.array([-0.2]),
    du_max=np.array([0.2]),
)

# Create MPC solver
mpc = MPC(
    dynamics=dynamics,
    state_dim=state_dim,
    action_dim=action_dim,
    horizon=20,
    weights=weights,
    constraints=constraints,
    iters=60,
    lr=0.05,
    optimizer="adam",
    warm_start=True,
)

# Solve
x0 = np.zeros(state_dim)
x_ref = np.zeros((21, state_dim))  # horizon+1 reference trajectory
x_ref[:, -1] = 0.1  # target for last state component

result = mpc.solve(x0=x0, x_ref=x_ref, u_prev=None)
print("First control:", result.u0)
print("Predicted trajectory shape:", result.x_seq.shape)

MPCAgent with learned dynamics (recommended)¶

MPCAgent provides a complete workflow: data collection, dynamics training, and MPC control.

import gymnasium as gym
import numpy as np
from tensoraerospace.agent.mpc import (
    MPCAgent,
    MPCWeights,
    MPCConstraints,
    MPCStepResponseExtraCostConfig,
)

# Create environment
env = gym.make("LinearLongitudinalB747-v0", ...)

# Configure weights
weights = MPCWeights(
    Q_diag=np.array([1.0, 1.0, 10.0, 100.0]),
    R_diag=np.array([0.01]),
    S_diag=np.array([0.5]),
    terminal_weight=1.0,
)

# Configure constraints
constraints = MPCConstraints(
    u_min=np.array([-0.3]),
    u_max=np.array([0.3]),
    du_min=np.array([-0.05]),
    du_max=np.array([0.05]),
)

# Extra cost for step response quality
step_cfg = MPCStepResponseExtraCostConfig.from_degrees(
    tracked_idx=-1,        # last state = theta
    rate_idx=-2,           # second-to-last = q (pitch rate)
    dt=0.01,
    overshoot_limit_deg=0.05,
    settle_band_deg=0.1,
    settle_time_target_s=1.0,
)

# Create agent
agent = MPCAgent(
    env,
    horizon=30,
    weights=weights,
    constraints=constraints,
    tracking_type="step_response",
    step_response_config=step_cfg,
    hidden_layers=(256, 256),
    normalize=True,
    device="cuda",  # or "cpu"
)

# Collect training data
agent.collect_data(num_episodes=50, exploration="signals")

# Train dynamics model
agent.train_dynamics(epochs=10, batch_size=1024)

# Use in control loop
obs, info = env.reset()
state = ...  # extract internal state from env
x_ref = ...  # reference trajectory (horizon+1, state_dim)

action = agent.select_action(state, x_ref=x_ref)
obs, reward, done, truncated, info = env.step(action)

# Save/load checkpoints
path = agent.save("./runs")
agent.load(path)

Using custom dynamics models¶

You can plug in different neural network architectures:

MLP (default)NARXTransformer

from tensoraerospace.agent.mpc import OneStepMLP

model = OneStepMLP(
    input_dim=state_dim + action_dim,
    output_dim=state_dim,
    hidden_layers=(256, 256, 128),
    activation="relu",  # or "tanh", "gelu"
)

agent = MPCAgent(env, model=model, ...)

from tensoraerospace.agent.mpc import NARXDynamicsModel

model = NARXDynamicsModel(
    state_dim=state_dim,
    action_dim=action_dim,
    hidden_size=256,
    num_layers=3,
    state_lags=1,
    control_lags=1,
)

agent = MPCAgent(env, model=model, ...)

from tensoraerospace.agent.mpc import TransformerDynamicsModel

model = TransformerDynamicsModel(
    input_dim=state_dim + action_dim,
    output_dim=state_dim,
    d_model=64,
    nhead=4,
    num_encoder_layers=2,
    dim_feedforward=256,
    dropout=0.1,
)

agent = MPCAgent(env, model=model, ...)

Extra Cost Functions¶

Tracking mode (`tracking`)¶

Adds penalties for control smoothness:

w_du: weight for \(\sum (\Delta u)^2\)
w_jerk: weight for \(\sum (\Delta^2 u)^2\)

from tensoraerospace.agent.mpc import MPCTrackingExtraCostConfig

cfg = MPCTrackingExtraCostConfig(w_du=50.0, w_jerk=10.0)
agent = MPCAgent(env, tracking_type="tracking", tracking_config=cfg, ...)

Step response mode (`step_response`)¶

Adds penalties for overshoot, settling time, oscillations:

from tensoraerospace.agent.mpc import MPCStepResponseExtraCostConfig

cfg = MPCStepResponseExtraCostConfig.from_degrees(
    tracked_idx=-1,               # index of tracked state (e.g., theta)
    rate_idx=-2,                  # index of rate state (e.g., q)
    dt=0.01,                      # timestep
    overshoot_limit_deg=0.05,     # max overshoot in degrees
    settle_band_deg=0.10,         # settling band width
    settle_time_target_s=1.0,     # target settling time
    w_overshoot=8000.0,           # overshoot penalty weight
    w_settle=8000.0,              # settling penalty weight
    w_osc=500.0,                  # oscillation penalty weight
    w_jerk=50.0,                  # jerk penalty weight
)

agent = MPCAgent(env, tracking_type="step_response", step_response_config=cfg, ...)

You can switch modes at runtime:

agent.set_tracking_type("tracking", tracking_config=tracking_cfg)
# or
agent.set_tracking_type("step_response", step_response_config=step_cfg)

Data Collection¶

MPCAgent.collect_data() supports two exploration strategies:

Strategy	Description
`"random"`	Random actions from `env.action_space.sample()`
`"signals"`	Rich signal library: steps, ramps, sinusoids, chirps, doublets, etc.

agent.collect_data(
    num_episodes=50,
    max_steps=1000,
    exploration="signals",
    signal_kinds=["random_steps", "sinusoid", "chirp", "doublet"],
    action_amplitude_frac=0.8,
)

Available signal types: random_steps, unit_step, multi_step, ramp, sinusoid, multisine, chirp, square_wave, triangular_wave, sawtooth, doublet, pulse, gaussian_pulse, exponential, damped_sinusoid.

Hyperparameters¶

MPC Solver (`MPC`)¶

Parameter	Description	Default
`horizon`	Prediction horizon	20
`iters`	Optimization iterations per solve	60
`lr`	Learning rate	0.05
`optimizer`	`"adam"` or `"sgd"`	`"adam"`
`warm_start`	Reuse previous solution	`True`
`track_best`	Track best solution during optimization	`True`
`compile_dynamics`	Use `torch.compile` (PyTorch 2.x)	`False`

MPCAgent¶

Parameter	Description	Default
`hidden_layers`	MLP hidden layer sizes	`(256, 256)`
`normalize`	Normalize inputs/outputs	`True`
`dynamics_lr`	Learning rate for dynamics model	`1e-3`
`grad_clip_norm`	Gradient clipping	`1.0`
`memory_capacity`	Replay buffer size	`200_000`
`model_predict_delta`	Predict \(\Delta x\) instead of \(x'\)	`True`

Best practices

Use exploration="signals" for better coverage of state-action space
Start with horizon=20-30 and increase if needed
Enable normalize=True for neural dynamics
Use tracking_type="step_response" for aerospace control tasks
For real-time control, consider compile_dynamics=True on GPU

Examples¶

Complete end-to-end examples demonstrating MPC with different dynamics models on the B747 longitudinal control task:

Example	Dynamics Model	Description
MPC + MLP	`OneStepMLP`	Standard MLP-based dynamics learning with step response tracking
MPC + NARX	`NARXDynamicsModel`	Nonlinear autoregressive model with exogenous inputs
MPC + Transformer	`TransformerDynamicsModel`	Transformer encoder for dynamics prediction

Each example demonstrates the full pipeline:

Environment setup — Create B747 environment with step reference signal for pitch (θ)
Data collection — Collect transitions using rich exploration signals
Dynamics training — Train neural network to predict state transitions
MPC rollout — Run closed-loop control using learned dynamics
Evaluation — Analyze step response quality (overshoot, settling time, etc.)

Key results from examples¶

Model	Overshoot	Settling Time	Rise Time	Static Error
MLP	~0.30%	~1.7s	~1.1s	~0.001
NARX	~-1.9%	~3.0s	~1.0s	~0.026
Transformer	~-0.10%	~1.5s	~0.8s	~0.009

Running examples

Examples are Jupyter notebooks located in example/mpc_controllers/. Run them to see full training logs, plots, and benchmark reports.

API Reference¶

`MPC(*, dynamics, state_dim, action_dim, horizon=20, weights, constraints=None, extra_cost_fn=None, iters=60, lr=0.05, optimizer='adam', device=None, dtype=torch.float32, warm_start=True, track_best=True, best_check_every=1, compile_dynamics=False, compile_mode='reduce-overhead', seed=None)` ¶

Projected-gradient MPC over a differentiable dynamics model.

This solver optimizes a control sequence U using torch/autograd and applies hard constraints via projection (clamp + sequential rate limiting).

`reset()` ¶

Reset warm-start state.

`solve(*, x0, x_ref=None, u_prev=None)` ¶

Solve MPC problem for the current state.

Parameters:

Name	Type	Description	Default
`x0`	`TensorLike`	Current state, shape (state_dim,) or (state_dim, 1) or (1, state_dim).	required
`x_ref`	`TensorLike \| None`	Optional reference trajectory for states. Expected shape is (horizon+1, state_dim) or (horizon, state_dim). If provided as (horizon, state_dim), it is interpreted as targets for x_{t+1} and a terminal target is appended by repeating the last row.	`None`
`u_prev`	`TensorLike \| None`	Optional previous control input, shape (action_dim,).	`None`

Returns:

Name	Type	Description
`MPCSolveResult`	`MPCSolveResult`	contains the first control input and the predicted trajectory.

MPCAgent(env, *, state_dim=None, action_dim=None, horizon=20, weights=None, constraints=None, tracking_type='tracking', tracking_config=None, step_response_config=None, extra_cost_fn=None, iters=60, mpc_lr=0.05, mpc_optimizer='adam', warm_start=True, mpc_track_best=True, mpc_best_check_every=1, mpc_compile_dynamics=False, mpc_compile_mode='reduce-overhead', model=None, model_predict_delta=True, hidden_layers=(256, 256), activation='relu', normalize=True, dynamics_lr=0.001, weight_decay=0.0, grad_clip_norm=1.0, memory_capacity=200000, obs_to_state=None, action_to_env=None, action_from_env=None, device='cpu', dtype=torch.float32, seed=0) ¶

Bases: BaseRLModel

DSAC-like wrapper around MPC with learned dynamics.

Goals: - Work with different Gymnasium-like environments (infer dims and bounds) - Accept a neural dynamics model at init (or build a default MLP) - Provide DSAC-ish ergonomics: buffer, collect_data(), train_dynamics(), select_action()

`set_tracking_type(tracking_type, *, tracking_config=None, step_response_config=None)` ¶

Switch extra-cost mode (tracking vs step_response).

`to_device(device)` ¶

Move model, normalizers, and MPC to a new device (DSAC-style).

`get_env()` ¶

Return current env.

`get_param_env()` ¶

Return env/policy metadata for HuggingFace-style checkpoints.

`reset()` ¶

Reset MPC warm-start and previous action memory.

`select_action(state, *, x_ref=None)` ¶

Compute action (env units) using MPC over learned dynamics.

`collect_data(*, num_episodes=10, max_steps=None, exploration='random', signal_kinds=None, dt=None, action_amplitude_frac=0.8)` ¶

Collect (x, u, x_next) transitions into self.memory.

Notes

For exploration="random" uses env.action_space.sample().
For exploration="signals" uses tensoraerospace.signals to generate time-series actions (works for continuous Box actions).

`fit_normalizers(*, num_samples=50000)` ¶

Fit x/u/y normalizers from a random subset of the replay buffer.

`train_dynamics(*, epochs=5, batch_size=1024, steps_per_epoch=None, loss='mse', force_refit_normalizers=False)` ¶

Train the dynamics model on transitions stored in the replay buffer.

The model is trained on samples from self.memory.

Parameters:

Name	Type	Description	Default
`force_refit_normalizers`	`bool`	If True, re-fit the normalizers even if they have already been fitted. Default is False — normalizers are fit only on the first call to avoid changing the loss landscape mid-training.	`False`

`save(path=None, save_gradients=True)` ¶

Save MPC agent in HuggingFace-style layout (config + weights).

`from_pretrained(repo_name, access_token=None, version=None, load_gradients=False)` `classmethod` ¶

Load checkpoint from local dir or HuggingFace Hub.

`push_to_hub(repo_name, access_token=None, save_path=None, include_gradients=False)` ¶

Save checkpoint and upload it to HuggingFace Hub.

`MPCWeights(Q_diag, R_diag, S_diag=None, terminal_weight=1.0)` `dataclass` ¶

Quadratic weights for the standard MPC objective.

The objective is

sum_{t=0..N-1} ||x_{t+1} - x_ref_{t+1}||Q^2 + ||u_t||_R^2 + ||u_t - u||_S^2 + terminal_weight * ||x_N - x_ref_N||_Q^2

Q, R, S are interpreted as diagonal weights (vectors). This keeps the API simple and fast enough for small horizons without extra dependencies.

`MPCConstraints(u_min=None, u_max=None, du_min=None, du_max=None)` `dataclass` ¶

Box constraints for control and rate limits.

All bounds are interpreted element-wise (per control dimension).

`MPCSolveResult(u0, u_seq, x_seq, final_cost, iters)` `dataclass` ¶

Result bundle for MPC solve.

`MPCTrackingExtraCostConfig(w_du=0.0, w_jerk=0.0)` `dataclass` ¶

Extra cost for generic reference tracking.

This is applied in addition to the quadratic MPC objective. Values are dimensionless weights.

`MPCStepResponseExtraCostConfig(tracked_idx=-1, rate_idx=None, dt=0.1, ref_change_threshold=float(np.deg2rad(0.1)), min_step_amp=float(np.deg2rad(0.5)), overshoot_limit=float(np.deg2rad(0.05)), settle_band=float(np.deg2rad(0.1)), settle_band_min=float(np.deg2rad(0.05)), settle_band_ratio=0.01, settle_time_target_s=1.0, rate_settle=float(np.deg2rad(0.25)), w_overshoot=8000.0, w_time=800.0, w_settle=8000.0, w_sse_steady=40000.0, w_osc=500.0, w_jerk=50.0, w_du_steady=80.0, w_jerk_steady=800.0)` `dataclass` ¶

Extra cost tuned for step response (overshoot/settling/osc/jerk).

All thresholds must be in the SAME UNITS as the tracked state component inside x_seq/x_ref (for B747 internal model this is radians).

`MPCStandardScaler(mean, std)` `dataclass` ¶

Simple per-feature standardization helper (mean/std).

The scaler lives on a torch device and is used both for training and for MPC rollouts through a learned dynamics model.

`OneStepMLP(*, input_dim, output_dim, hidden_layers=(256, 256), activation='relu')` ¶

Bases: Module

A small MLP for one-step dynamics learning.

Expected IO

in: concatenated [x, u] of shape (B, state_dim + action_dim) out: either delta-x or x_next of shape (B, state_dim)

`forward(xu)` ¶

Forward pass.

`NARXDynamicsModel(*, state_dim, action_dim, hidden_size=256, num_layers=2, state_lags=1, control_lags=1)` ¶

Bases: Module

MPCAgent-compatible NARX dynamics model.

MPCAgent expects learned dynamics modules with signature:

y = model(xu), where xu = concat([x, u])

This wrapper builds an internal :class:~tensoraerospace.agent.mpc.narx.NARX and provides the required forward(xu) interface.

Notes

For the current MPC pipeline, this is typically used with state_lags=1 and control_lags=1 (one-step model).
If you want true NARX with lags > 1, you must provide an augmented state/action history vector as input to MPC (not handled implicitly).

`forward(xu)` ¶

Forward pass.

Parameters:

Name	Type	Description	Default
`xu`	`Tensor`	Concatenated input of shape (B, state_dimstate_lags + action_dimcontrol_lags).	required

Returns:

Type	Description
`Tensor`	Predicted next-state (or delta-state) of shape (B, state_dim).

`NARX(input_size, hidden_size, output_size, num_layers, state_lags, control_lags)` ¶

Bases: Module

NARX neural network for learning system dynamics.

The model uses lagged (historical) state and control inputs to predict the next state.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Input size (concatenated lagged states and controls).	required
`hidden_size`	`int`	Hidden layer size.	required
`output_size`	`int`	Output size (predicted state dimension).	required
`num_layers`	`int`	Number of hidden layers.	required
`state_lags`	`int`	Number of state lags used as input.	required
`control_lags`	`int`	Number of control lags used as input.	required

Initialize the NARX network.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Input size.	required
`hidden_size`	`int`	Hidden layer size.	required
`output_size`	`int`	Output size.	required
`num_layers`	`int`	Number of hidden layers.	required
`state_lags`	`int`	Number of state lags.	required
`control_lags`	`int`	Number of control lags.	required

`forward(state, control)` ¶

Run a forward pass.

Parameters:

Name	Type	Description	Default
`state`	`Tensor`	Lagged state tensor.	required
`control`	`Tensor`	Lagged control tensor.	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Predicted next state.

`TransformerDynamicsModel(input_dim, output_dim, d_model=64, nhead=4, num_encoder_layers=2, dim_feedforward=256, dropout=0.1, seq_len=1)` ¶

Bases: Module

Transformer-based model for learning system dynamics.

Predicts the next system state from a sequence of state+action inputs using a Transformer encoder.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Input dimension (state + control).	required
`output_dim`	`int`	Output dimension (next state).	required
`d_model`	`int`	Transformer model dimension. Defaults to `64`.	`64`
`nhead`	`int`	Number of attention heads. Defaults to `4`.	`4`
`num_encoder_layers`	`int`	Number of encoder layers. Defaults to `2`.	`2`
`dim_feedforward`	`int`	Feed-forward layer dimension. Defaults to `256`.	`256`
`dropout`	`float`	Dropout probability. Defaults to `0.1`.	`0.1`
`seq_len`	`int`	Sequence length. Defaults to `1`.	`1`

Initialize the transformer dynamics model.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Input dimension.	required
`output_dim`	`int`	Output dimension.	required
`d_model`	`int`	Model dimension.	`64`
`nhead`	`int`	Number of attention heads.	`4`
`num_encoder_layers`	`int`	Number of encoder layers.	`2`
`dim_feedforward`	`int`	Feed-forward dimension.	`256`
`dropout`	`float`	Dropout probability.	`0.1`
`seq_len`	`int`	Sequence length.	`1`

`forward(x)` ¶

Forward pass through transformer model.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor of shape (batch_size, input_dim).	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Predicted next state of shape (batch_size, output_dim).

Model Predictive Control (MPC)¶

Theory (brief)¶

Architecture¶

Quick Start¶

Basic MPC with a custom dynamics function¶

MPCAgent with learned dynamics (recommended)¶

Using custom dynamics models¶

Extra Cost Functions¶

Tracking mode (tracking)¶

Step response mode (step_response)¶

Data Collection¶

Hyperparameters¶

MPC Solver (MPC)¶

MPCAgent¶

Examples¶

Key results from examples¶

API Reference¶

MPC(*, dynamics, state_dim, action_dim, horizon=20, weights, constraints=None, extra_cost_fn=None, iters=60, lr=0.05, optimizer='adam', device=None, dtype=torch.float32, warm_start=True, track_best=True, best_check_every=1, compile_dynamics=False, compile_mode='reduce-overhead', seed=None) ¶

reset() ¶

solve(*, x0, x_ref=None, u_prev=None) ¶

set_tracking_type(tracking_type, *, tracking_config=None, step_response_config=None) ¶

to_device(device) ¶

get_env() ¶

get_param_env() ¶

reset() ¶

select_action(state, *, x_ref=None) ¶

collect_data(*, num_episodes=10, max_steps=None, exploration='random', signal_kinds=None, dt=None, action_amplitude_frac=0.8) ¶

fit_normalizers(*, num_samples=50000) ¶

train_dynamics(*, epochs=5, batch_size=1024, steps_per_epoch=None, loss='mse', force_refit_normalizers=False) ¶

save(path=None, save_gradients=True) ¶

from_pretrained(repo_name, access_token=None, version=None, load_gradients=False) classmethod ¶

push_to_hub(repo_name, access_token=None, save_path=None, include_gradients=False) ¶

MPCWeights(Q_diag, R_diag, S_diag=None, terminal_weight=1.0) dataclass ¶

MPCConstraints(u_min=None, u_max=None, du_min=None, du_max=None) dataclass ¶

MPCSolveResult(u0, u_seq, x_seq, final_cost, iters) dataclass ¶

MPCTrackingExtraCostConfig(w_du=0.0, w_jerk=0.0) dataclass ¶

MPCStandardScaler(mean, std) dataclass ¶

OneStepMLP(*, input_dim, output_dim, hidden_layers=(256, 256), activation='relu') ¶

forward(xu) ¶

NARXDynamicsModel(*, state_dim, action_dim, hidden_size=256, num_layers=2, state_lags=1, control_lags=1) ¶

forward(xu) ¶

NARX(input_size, hidden_size, output_size, num_layers, state_lags, control_lags) ¶

forward(state, control) ¶

TransformerDynamicsModel(input_dim, output_dim, d_model=64, nhead=4, num_encoder_layers=2, dim_feedforward=256, dropout=0.1, seq_len=1) ¶

forward(x) ¶

Tracking mode (`tracking`)¶

Step response mode (`step_response`)¶

MPC Solver (`MPC`)¶

`MPC(*, dynamics, state_dim, action_dim, horizon=20, weights, constraints=None, extra_cost_fn=None, iters=60, lr=0.05, optimizer='adam', device=None, dtype=torch.float32, warm_start=True, track_best=True, best_check_every=1, compile_dynamics=False, compile_mode='reduce-overhead', seed=None)` ¶

`reset()` ¶

`solve(*, x0, x_ref=None, u_prev=None)` ¶

`set_tracking_type(tracking_type, *, tracking_config=None, step_response_config=None)` ¶

`to_device(device)` ¶

`get_env()` ¶

`get_param_env()` ¶

`reset()` ¶

`select_action(state, *, x_ref=None)` ¶

`collect_data(*, num_episodes=10, max_steps=None, exploration='random', signal_kinds=None, dt=None, action_amplitude_frac=0.8)` ¶

`fit_normalizers(*, num_samples=50000)` ¶

`train_dynamics(*, epochs=5, batch_size=1024, steps_per_epoch=None, loss='mse', force_refit_normalizers=False)` ¶

`save(path=None, save_gradients=True)` ¶

`from_pretrained(repo_name, access_token=None, version=None, load_gradients=False)` `classmethod` ¶

`push_to_hub(repo_name, access_token=None, save_path=None, include_gradients=False)` ¶

`MPCWeights(Q_diag, R_diag, S_diag=None, terminal_weight=1.0)` `dataclass` ¶

`MPCConstraints(u_min=None, u_max=None, du_min=None, du_max=None)` `dataclass` ¶

`MPCSolveResult(u0, u_seq, x_seq, final_cost, iters)` `dataclass` ¶

`MPCTrackingExtraCostConfig(w_du=0.0, w_jerk=0.0)` `dataclass` ¶

`MPCStandardScaler(mean, std)` `dataclass` ¶

`OneStepMLP(*, input_dim, output_dim, hidden_layers=(256, 256), activation='relu')` ¶

`forward(xu)` ¶

`NARXDynamicsModel(*, state_dim, action_dim, hidden_size=256, num_layers=2, state_lags=1, control_lags=1)` ¶

`forward(xu)` ¶

`NARX(input_size, hidden_size, output_size, num_layers, state_lags, control_lags)` ¶

`forward(state, control)` ¶

`TransformerDynamicsModel(input_dim, output_dim, d_model=64, nhead=4, num_encoder_layers=2, dim_feedforward=256, dropout=0.1, seq_len=1)` ¶

`forward(x)` ¶