Co-state Neural Network for Real-time Nonlinear Optimal Control with Input Constraints

TL;DR

This work addresses real-time solution of nonlinear OCPs with control input constraints by introducing a co-state neural network (CoNN) that learns the mapping from an initial state $x_0$ to its optimal co-state trajectory under Pontryagin's Minimum Principle (PMP). The CoNN parameterizes the TPBVP, and online control is obtained by solving a QP using the predicted co-states to enforce input constraints, yielding an MPC-like closed-loop controller with reduced computational burden. Key contributions include a novel training framework with two loss components (prediction and continuity) to enforce PMP-consistent co-states, and demonstration on a 1D nonlinear system showing robustness to unseen initial conditions and disturbances, with results comparable to direct solvers. The approach promises scalable real-time constrained OCP solutions and provides a pathway to extending to higher-dimensional problems and explicit state-constraint handling in training.

Abstract

In this paper, we propose a method to solve nonlinear optimal control problems (OCPs) with constrained control input in real-time using neural networks (NNs). We introduce what we have termed co-state Neural Network (CoNN) that learns the mapping from any given state value to its corresponding optimal co-state trajectory based on the Pontryagin's Minimum (Maximum) Principle (PMP). In essence, the CoNN parameterizes the Two-Point Boundary Value Problem (TPBVP) that results from the PMP for various initial states. The CoNN is trained using data generated from numerical solutions of TPBVPs for unconstrained OCPs to learn the mapping from a state to its corresponding optimal co-state trajectory. For better generalizability, the CoNN is also trained to respect the first-order optimality conditions (system dynamics). The control input constraints are satisfied by solving a quadratic program (QP) given the predicted optimal co-states. We demonstrate the effectiveness of our CoNN-based controller in a feedback scheme for numerical examples with both unconstrained and constrained control input. We also verify that the controller can handle unknown disturbances effectively.

Figures (5)

• Figure 1: CoNN architecture and training procedures illustrated for $x_0^{(i)} \in \mathbb{R}, \forall i$, same as the one-dimensional system shown in the Sec. V. Note that the index [$j$] denote the $j^{th}$ entry of the trajectory starting from 0.
• Figure 2: CoNN-based controller block diagram for validation.
• Figure 3: CoNN-based controller with unseen initial conditions.
• Figure 4: CoNN-based controller with control input constraints
• Figure 5: CoNN-based controller with disturbance