Finite-Horizon Optimal Control Problems

The models introduced in the previous chapter predict how actions change a system’s state. Trajectory optimization adds an objective and constraints, then selects the actions that produce a desirable state sequence. For a fixed discrete horizon, the states and controls form a finite vector, so the planning problem can be written as a nonlinear program.

Which finite-dimensional optimization problem selects the best admissible trajectory from one known initial condition?

A trajectory is the time-indexed sequence of states $(\mathbf{x}_1,\ldots,\mathbf{x}_T)$ and controls $(\mathbf{u}_1,\ldots,\mathbf{u}_{T-1})$. This chapter first computes the complete control sequence from a known initial state. The sequence is open loop: once execution begins, the action at time $t$ does not change in response to the measured state.

The three-satellite example below makes both the value and the limitation of open-loop planning visible. A linear model produces a sparse, feasible plan, but the same immutable plan misses its target when replayed through a nonlinear model. Later chapters will replace the fixed sequence by feedback through receding-horizon control and policies.

A Motivating Example: Phasing Three Satellites with Differential Drag¶

What does an open-loop plan gain from a transparent finite-horizon model, and which part of its guarantee disappears when the replay model changes?

Three small satellites are released into nearly the same circular orbit, but the mission calls for them to occupy slots separated by $120^\circ$. They have no propulsion. Each satellite can instead rotate between low- and high-drag attitudes. High drag lowers its orbit slightly; the lower satellite then moves faster and accumulates phase relative to the others. The control authority is weak, slow, and irreversible because every maneuver spends altitude.

Differential drag has been used to phase propulsionless satellite constellations in orbit Foster et al., 2018. Linear programs provide a useful planning model for this control mechanism Sin et al., 2018. The example below is a teaching-scale reconstruction inspired by that literature. It is not a reconstruction of a particular flight campaign.

We plan over $N=180$ daily intervals. All three satellites begin at a circular altitude of 475 km, with phase offsets

and zero relative angular rates. The action

is the fraction of day $k$ that satellite $i$ spends in its high-drag attitude. For the nominal daily model, its state is

where $\varphi$ is phase in degrees, $\omega$ is relative angular rate in degrees per day, and $\ell$ is extra altitude loss in kilometres relative to remaining in the low-drag attitude.

From drag physics to a daily linear model¶

Use SI units in the derivation. With $h_0=475\text{ km}$, set

At the reference density $\rho_0$, changing the ballistic coefficient from $B_\mathrm{low}$ to $B_\mathrm{high}$ changes the area-to-mass factor by

Linearizing the semimajor-axis and mean-motion changes over one day gives

where $\Delta t=86400$ s. For

the coefficients are

The daily dynamics are therefore

The phase receives half of the new daily rate during the interval, the rate accumulates the drag-induced acceleration, and the extra altitude loss accumulates monotonically.

A finite open-loop plan¶

Define the cyclic difference matrix

The terminal target is

with tolerances of 0.1 degree and 0.002 degree per day. These are unwrapped directed differences. Because the rows of $G$ sum to zero, the third target is $-240^\circ$, which represents the same circular separation as $+120^\circ$ modulo $360^\circ$ while preserving a consistent unwrapped coordinate system.

The primary linear program minimizes the worst final extra altitude loss:

A second linear program keeps $z$ at its primary optimum, up to numerical tolerance, and minimizes

This lexicographic step selects a low-variation member of the primary optimal set without changing the worst-loss objective beyond numerical tolerance.

Nominal and nonlinear replay¶

The linear plan is first rolled out through the model used by the optimizer. The exact same $u_{i,k}$ is then replayed, without reoptimization, through the nonlinear orbital model

where

Here $a_i$ and $R_\mathrm{E}$ are in metres inside the orbital equations, so $h_i$ is the corresponding altitude in kilometres. In the display, the nominal altitude trace is the planning-model proxy $475-\ell_i$ km and therefore omits the common low-drag decay; the nonlinear trace reports absolute orbital altitude.

The nonlinear trajectory is integrated by hourly RK4 and checked against a 30-minute replay. This variable-density model is a deterministic teaching stress test, not a flight-dynamics reconstruction. The complete command sequence is known from day zero; the state traces below are revealed only up to the playhead.

from pathlib import Path
import sys

from IPython.display import HTML, display

code_dir = Path.cwd() / "code"
if str(code_dir) not in sys.path:
    sys.path.insert(0, str(code_dir))

from cubesat_replay import render_cubesat_replay

display(HTML(render_cubesat_replay(
    Path("artifacts/cubesat/textbook_results.json"),
    replay_id="cubesat-differential-drag-replay",
    fallback_id="fig-cubesat-differential-drag-fallback",
)))

The lexicographic solve limits the largest nominal extra altitude loss to 1.312323 km, then reduces total variation from 33.701181 to 2.000000 without changing that primary answer beyond the declared lock tolerance.

At day 180, the linear planning model reaches a maximum cyclic-gap error of $0.100000^\circ$ and a maximum cyclic relative rate of $0.002000^\circ/\mathrm{day}$. Replaying the identical command matrix through the variable-density model raises those values to $11.537412^\circ$ and $0.031144^\circ/\mathrm{day}$. The optimization result is therefore a statement about the planning model, not a guarantee for the replay model.

The hourly RK4 replay agrees with a 30-minute reference to within $3.101e-09^\circ$ in phase, $6.333e-11^\circ/\mathrm{day}$ in relative rate, and 3.183e-11 km in altitude. Its sampled density spans 2.264e-13 to 3.403e-13 kg/m³.

All 15 deterministic acceptance checks pass. This is a transparent stress test of one open-loop teaching model, not a flight-operations prescription.

Download the open-loop plan (CSV)

Download the audit metrics (CSV)

The nominal rollout establishes feasibility for the optimization model. The nonlinear replay tests the same plan under the declared model change, and the terminal phase constraints fail. Closing the loop by replanning will replace the immutable schedule by controls that can change when new state measurements arrive.

The example already contains the ingredients of a discrete-time optimal control problem (DOCP): a state $x_{i,k}$, a bounded control $u_{i,k}$, a transition map, terminal constraints, and an objective accumulated over a finite horizon. We now formalize that structure.

Discrete-Time Optimal Control Problems (DOCPs)¶

Which variables, costs, dynamics, and constraints place the satellite planner inside a reusable finite-horizon control template?

Consider a system described by a state $\mathbf{x}_t \in \mathbb{R}^n$, summarizing everything needed to predict its evolution. At each stage $t$, we can influence the system through a control input $\mathbf{u}_t \in \mathbb{R}^m$. The dynamics specify how the state evolves:

where $\mathbf{f}_t$ may be nonlinear or time-varying. We assume the initial state $\mathbf{x}_1$ is known.

The goal is to pick a sequence of controls $\mathbf{u}_1,\dots,\mathbf{u}_{T-1}$ that makes the trajectory desirable. But desirable in what sense? That depends on an objective function, which often includes two components:

The stage cost reflects ongoing penalties such as energy, delay, or risk. The terminal cost measures the value (or cost) of ending in a particular state. Together, these give a discrete-time Bolza problem with path constraints and bounds:

In the satellite example, $\mathbf{x}_t$ stacks the three phase, relative-rate, and altitude-loss states; $\mathbf{u}_t$ stacks the three daily drag fractions; and $\mathbf{f}_t$ applies the block-diagonal copies of the daily map. The terminal inequalities impose the cyclic slot and rate tolerances, while the epigraph variable $z$ represents the worst altitude loss. This mapping also shows why state values can appear explicitly as decision variables even when a deterministic rollout could reconstruct them from the controls.

Written this way, it may seem obvious that the decision variables are the controls $\mathbf{u}_t$. After all, in most intuitive descriptions of control, we think of choosing inputs to influence the system. But notice that in the program above, the entire state trajectory also appears as a set of variables, linked to the controls by the dynamics constraints. This is intentional: it reflects one way of writing the problem that makes the constraints explicit.

Why introduce $\mathbf{x}_t$ as decision variables if they can be simulated forward from the controls? Many readers hesitate here, and the question is natural: If the model is deterministic and $\mathbf{x}_1$ is known, why not pick $\mathbf{u}_{1:T-1}$ and compute $\mathbf{x}_{2:T}$ on the fly? That instinct leads to single shooting, a method we will return to shortly.

Already in this formulation, though, the structure of the problem matters. Ignoring it can make our life much harder. The reason is twofold:

• Dimensionality grows with the horizon. For a horizon of length $T$, the program has roughly $(T-1)(m+n)$ decision variables.

• Temporal coupling. Each control affects all future states and costs. The feasible set is not a simple box but a narrow manifold defined by the dynamics.

Together, these features explain why specialized methods exist and why the way we write the problem influences the algorithms we can use. Whether we keep states explicit or eliminate them through forward simulation determines the problem size, its conditioning, and the trade-offs between robustness and computational effort.

Existence of Solutions and Optimality Conditions¶

Writing the problem as a nonlinear program does not ensure that a minimizer exists or that a candidate is locally optimal. Which conditions supply those two claims?

Now that we have the optimization problem written down, we can ask: does it always have a solution? And if so, how do we recognize one? These questions lead us to feasibility and optimality conditions.

Existence of Solutions¶

Notice first that nothing in the problem statement required the dynamics

to be stable. In fact, many problems of interest involve unstable systems; think of balancing a pole or steering a spacecraft. What matters is that the dynamics are well defined: given a state–control pair, the rule $\mathbf{f}_t$ produces a valid next state.

In continuous time, one usually requires $\mathbf{f}$ to be continuous (often Lipschitz continuous) in $\mathbf{x}$ so that the ODE has a unique solution on the horizon of interest. In discrete time, the requirement is lighter: we only need the update map to be well posed.

Existence also hinges on feasibility. A candidate control sequence must generate a trajectory that respects all constraints: the dynamics, any bounds on state and control, and any terminal requirements. If no such sequence exists, the feasible set is empty and the problem has no solution. This can happen if the constraints are overly strict, or if the system is uncontrollable from the given initial condition.

Optimality Conditions¶

Assume the feasible set is nonempty. Necessary conditions for a local minimum follow from the Lagrange multiplier theorem. They supply equations for candidate solutions, although satisfying them alone need not establish a minimum. For a smooth problem

Here $G$ stacks inequalities of the form $g_i\leq0$, and $H$ stacks equalities of the form $h_i=0$. Define the Lagrangian

The positive sign before the inequality term follows from the convention $G(\mathbf z)\le \mathbf0$: a violation $G_i(\mathbf z)>0$ increases the Lagrangian when $\mu_i>0$. Equality multipliers remain unrestricted. For a candidate point $\mathbf z$, the active set is

while indices with $G_i(\mathbf{z})<0$ are inactive. Only active inequalities can carry positive multipliers.

A constraint qualification is a regularity assumption under which a local minimum admits Lagrange multipliers. One sufficient choice is the Linear Independence Constraint Qualification (LICQ): the equality gradients and active inequality gradients at the candidate are linearly independent. If $J_H$ and $J_{G_{\mathcal A}}$ have those gradients as rows, LICQ requires

This assumption guarantees multiplier existence at a local minimum. It does not by itself make the Newton system for the optimality conditions invertible; that also depends on the objective and constraint curvature.

Under such a constraint qualification, any local minimizer $\mathbf{z}^\star$ admits multipliers $(\boldsymbol{\lambda}^\star,\boldsymbol{\mu}^\star)$ that satisfy the Karush–Kuhn–Tucker (KKT) conditions:

Complementarity requires $\mu_i^\star=0$ when $G_i(\mathbf z^\star)<0$. At a tight constraint, the multiplier may be positive or zero. Under suitable sensitivity assumptions, multipliers quantify changes in the optimal value when constraint right-hand sides change, with signs determined by the chosen perturbation convention.

In trajectory problems, $\mathbf z$ stacks states and controls, $H$ enforces the dynamics, and $G$ collects bounds and path constraints. The equality multipliers associated with dynamics become Pontryagin costates when the residual is written $\mathbf f_t-\mathbf x_{t+1}$. Writing the residual in the opposite order reverses the multiplier sign. Terminal-equality multipliers are denoted separately by $\boldsymbol\nu$ in that derivation. Since nonlinear dynamics generally make these programs nonconvex, KKT is a necessary first-order test under regularity. It does not generally imply a local minimum of the Lagrangian or a global saddle point.

If LICQ fails, a different constraint qualification may still apply. Without an applicable qualification, multipliers can fail to exist or be nonunique. Redundant dynamic or path constraints can also make numerical systems singular or poorly conditioned.

From KKT to algorithms¶

The Lagrangian gives an exact min–max representation of the constrained problem, whether or not the problem is convex:

Here $\boldsymbol\lambda$ is unrestricted. For a feasible $\mathbf z$, the inner supremum equals $F(\mathbf z)$. If an equality is violated, its multiplier can choose a sign and magnitude that makes the Lagrangian arbitrarily large. If $G_i(\mathbf z)>0$, sending $\mu_i$ to infinity does the same. Thus the multiplier player’s response assigns $+\infty$ to every infeasible choice.

Reversing the order gives the dual value, a lower bound on the primal value:

The distinction is whether the multipliers can respond to each candidate or must be fixed before the candidate is chosen. A saddle point requires one fixed multiplier pair for which $\mathbf z^\star$ globally minimizes the Lagrangian, together with optimal multipliers against that fixed $\mathbf z^\star$. Such a pair satisfies

for all $\mathbf z$, unrestricted $\boldsymbol\lambda$, and $\boldsymbol\mu\ge\mathbf0$. Solving the primal min–max problem alone does not guarantee this property. For example, minimizing $-z^2$ subject to $z=0$ has the unique feasible optimizer $z^\star=0$ and KKT multiplier $\lambda^\star=0$, yet $\mathcal L(z,0)=-z^2$ is maximized at zero.

When $F$ is convex, each $G_i$ is convex, and $H$ is affine, the Lagrangian is convex in $\mathbf z$. KKT stationarity then guarantees its global minimum; feasibility and complementarity give the other saddle inequality. An existing KKT point is therefore sufficient for a global saddle point. Slater’s condition, the existence of a point with $H=0$ and $G_i<0$ for every inequality, guarantees suitable multipliers at an attained finite optimizer in this convex setting. The nonlinear-programming appendix derives these statements and distinguishes them from the necessary conditions used in nonconvex trajectory optimization.

Primal–dual gradient dynamics (Arrow–Hurwicz)¶

A first-order method takes a descent step in the primal variables and an ascent step in the multipliers, projecting the inequality multipliers back onto their allowed set:

Here $\Pi_{\ge0}$ clips each negative component to zero. With the convention $G\le0$, a violated inequality increases its multiplier. These finite updates do not perform the inner supremum in the exact primal representation, and saddle-point existence alone does not guarantee their convergence. Even the convex bilinear case $\mathcal L(z,\lambda)=\lambda z$ has simultaneous updates with eigenvalues $1\pm i\sqrt{\alpha\beta}$ for fixed positive steps, whose magnitudes exceed one.

Additional curvature assumptions or algorithmic modifications are needed for convergence. Penalty and augmented-Lagrangian methods add terms that discourage violations. For example, one can use the penalized Lagrangian

where $\rho>0$ weights squared equality residuals and inequality violations. These terms change the curvature away from feasibility. Their use still requires a suitable step-selection or subproblem strategy; adding them alone is not a convergence guarantee.

SQP as Newton on the KKT system (equality case)¶

With only equality constraints $H(\mathbf{z})=\mathbf{0}$, write first-order conditions

Applying Newton’s method to this system gives the linear KKT solve

This is exactly the step computed by Sequential Quadratic Programming (SQP) in the equality-constrained case: it is Newton’s method on the KKT equations. For general problems with inequalities, SQP forms a quadratic subproblem by quadratically modeling $F$ with $\nabla_{\mathbf{z}\mathbf{z}}^2\mathcal{L}$ and linearizing the constraints, then solves that QP with line search or trust region. In least-squares-like problems one often uses Gauss–Newton (or a Levenberg–Marquardt trust region) as a positive-definite approximation to the Lagrangian Hessian.

In trajectory optimization, the KKT matrix inherits banded/sparse structure from the dynamics. Newton/SQP steps can be computed efficiently by exploiting this structure; in the special case of quadratic models and linearized dynamics, the QP reduces to an LQR solve along the horizon (this is the backbone of iLQR/DDP-style methods). Primal-dual updates provide simpler iterations and are easy to implement; augmented terms are typically needed to obtain stable progress when constraints couple stages.

The choice between methods depends on the context. Primal-dual gradients give lightweight iterations and are suited for warm starts or as inner loops with penalties. SQP/Newton can give rapid local convergence near a regular solution when the relevant second-order conditions also hold; trust regions or line search control steps farther from a solution.

Further Sources of Discrete-Time Optimal-Control Problems¶

Beyond sampled physical dynamics, which computations and continuous-time models produce the same temporally coupled optimization structure?

The satellite planner begins from a deliberately discretized daily model. Other problems are discrete because decisions naturally occur at stages, while still others inherit a discrete transition from numerical integration or program execution. The next two constructions make those latter connections explicit.

DOCPs Arising from the Discretization of Continuous-Time OCPs¶

Although many applications are natively discrete-time, it is also common to obtain a DOCP by discretizing a continuous-time formulation. Consider a system on $[0, T_c]$ given by

Choose a step size $\Delta > 0$ and grid $t_k = k\,\Delta$. A one-step integration scheme induces a discrete map $\mathbf{F}_\Delta$ so that

where, for example, explicit Euler gives $\mathbf{F}_\Delta(\mathbf{x},\mathbf{u},t) = \mathbf{x} + \Delta\,\mathbf{f}(\mathbf{x},\mathbf{u},t)$. The resulting discrete-time optimal control problem takes the Bolza form with these induced dynamics:

Programs as DOCPs and Differentiable Programming¶

It is often useful to view a computer program itself as a discrete-time dynamical system. Let the program state collect memory, buffers, and intermediate variables, and let the control represent inputs or tunable decisions at each step. A single execution step defines a transition map

and a scalar objective (e.g., loss, error, runtime, energy) yields a DOCP:

In differentiable programming (e.g., JAX, PyTorch), the composed map $\Phi_{T-1}\circ\cdots\circ\Phi_0$ is differentiable, enabling reverse-mode automatic differentiation and efficient gradient-based trajectory optimization. When parts of the program are non-differentiable (discrete branches, simulators with events), DOCPs can still be solved using derivative-free or weak-gradient methods (eg. finite differences, SPSA, Nelder–Mead, CMA-ES, or evolutionary strategies) optionally combined with smoothing, relaxations, or stochastic estimators to navigate non-smooth regions.

Example: Offline Frequency Planning for Inference¶

The inference service from the model-interface chapter is also a program with a controllable execution rate. Its service-rate and phase-power curves come from a measured NVIDIA L4 profile. The request and queue trajectories below come from a simulator calibrated with those curves, not from running each controller on the GPU. The scheduling rule is held fixed. Each 0.1-second simulator step prioritizes decode or begins with a prefill chunk capped at 512 tokens. If that chunk finishes early, the remaining service budget may return to decode. Active decode receives alternating-step or cache-pressure priority. This reduced interleaving model is not a reproduction of the vLLM scheduler used for profiling. The control sequence specifies one normalized GPU frequency for each second of a 60-second horizon,

An aggregate state collects queued prefill work, active decode work, temperature, and the preceding frequency:

The disturbance $w_k$ contains the arrival times and prompt lengths predicted for second $k$. Future output lengths remain hidden. The planner substitutes the trace distribution’s expected output length, while the request-level replay uses each realized length only as a disturbance. Service and power in $F_k$ interpolate the committed profile whose provenance is displayed with the result. The request-level simulator remains outside the optimizer and validates the resulting schedule after the solve.

The 60-second workload is selected before optimization by a deterministic, capacity-screened rule. The rule scans ten-second-aligned windows after load normalization, discards any window whose forecast work exceeds the horizon’s maximum-clock service capacity, and requires an occupied burst that can be moved twenty seconds earlier. Among the remaining windows, it chooses the shift-eligible burst with the largest forecast work; total window work and then earlier source time break ties. The selected source interval is $[890,950)$ seconds in the normalized trace and is rebased to start at zero. Its 48 requests require 51.936 seconds of forecast work at maximum clock, or 86.56% of the horizon’s capacity. The chosen burst occupies $[40,50)$ seconds after rebasing and contains 31 requests. The shifted replay moves those requests to $[20,30)$ seconds.

The offline problem uses the complete nominal arrival forecast. It introduces a normalized service decision $\nu_k\in[\nu_{\min},1]$ and a nonnegative backlog variable $B_{k+1}$. If $W_k$ is arriving work in seconds of highest-clock service, their fluid balance is relaxed to

The coefficient $\alpha_P$ is the slope of a linear interpolation between the lowest- and highest-clock normalized power values. The backlog term prices waiting throughout the horizon, while the additional terminal term discourages postponing work beyond second 60. This is a linear program, solved with HiGHS. The service decisions are mapped through the profiled service curve to continuous frequencies. Execution rounds each frequency downward to the nearest profiled requested clock, so the request-level validation includes the actuator’s finite action set. The replay reports the corresponding measured median realized clock separately; a requested level and its realized clock need not coincide under the experimental power cap.

This planning model deliberately omits request identities, phase-specific queues, clock slew, and the nonlinear thermal state. The detailed replay restores those variables and reports power, temperature, latency, and memory violations. The optimizer therefore supplies an offline plan from a tractable aggregate model, while the replay audits the assumptions used to obtain it.

The experiment asks what a one-shot frequency schedule gains from a perfect nominal arrival and prompt-length forecast, and what it loses when that forecast is wrong. Future output lengths remain uncertain in both cases. The same plan is replayed twice. The nominal replay uses the forecast supplied to the optimizer. The shifted replay uses the earlier arrival times defined above. No reoptimization occurs after either replay starts.

from pathlib import Path
import sys

from IPython.display import HTML, display

code_dir = Path.cwd() / "code"
if str(code_dir) not in sys.path:
    sys.path.insert(0, str(code_dir))

from inference_replay import render_serving_replay

display(HTML(render_serving_replay(
    Path("artifacts/inference_serving/textbook_results.json"),
    view="open_loop",
)))

The table reports request-level results for the fixed offline schedule under its nominal forecast and the shifted-burst disturbance. Both columns use the same requests, controller parameters, and measured profile calibration.

import pandas as pd

open_loop_metrics = pd.read_csv(
    "artifacts/inference_serving/metrics_open_loop.csv"
).set_index("controller")
open_loop_metrics[
    [
        "mean_ttft_s",
        "p95_ttft_s",
        "matched_moved_burst_mean_ttft_s",
        "matched_moved_burst_p95_ttft_s",
        "energy_j",
        "peak_queued_requests_at_minimum_clock",
        "queued_requests_at_30_s",
        "power_violation_w",
        "thermal_violation_c",
    ]
].T.rename_axis("metric").round(3)

Download every open-loop metric (CSV)

HiGHS reports an optimal solution for the stated linear program, with objective 11,205.68 in its weighted model units. Across all 48 requests, mean time to first token rises from 16.35 seconds under the nominal arrival times to 23.23 seconds after the shift. The 95th percentile rises from 28.09 to 31.28 seconds. For the 31 moved requests, the mean rises from 15.22 to 22.73 seconds. Their 95th percentile rises from 23.74 to 32.04 seconds. These changes are increases of 7.52 and 8.30 seconds, respectively.

At 30 seconds, the nominal replay has no queued request, while the shifted replay has 29. The peak queue while the clock is at its minimum also rises from zero to 29 requests. Energy falls from 3,719.7 to 3,646.8 joules despite the larger delays. Moving the burst changes which requests overlap and how long the system remains in each phase, so an energy decrease does not imply an improved service trajectory.

All requests eventually complete during the post-horizon drain. At the 60-second reporting horizon, 27 nominal requests and 22 shifted requests remain unfinished. Both simulations reach a modeled phase power of 64.852 W and exceed the configured 64.800 W power limit by 0.052 W. Neither simulation records a thermal or KV-capacity violation.

The displayed table focuses on latency, energy, queueing, and constraint violations. The downloadable CSV also reports energy per output token, time per output token, unfinished work, and the full set of recorded diagnostics. The nominal run tests the optimized trajectory under its own assumptions. The shifted run tests sensitivity to one explicit forecast error. It does not establish robustness to arbitrary arrivals, model error, or hardware throttling. Closing that gap requires new information to alter future controls, which is the role of feedback and receding-horizon optimization. The profile calibration is a hardware measurement, while the controller comparison is a simulation of the calibrated model. The reported latency, energy, and constraint outcomes are not direct measurements from replaying this trace through vLLM.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def optimize_open_loop(
    workload: Sequence[Request],
    plant: ServingPlant,
    scheduler: Scheduler = chunked_prefill_scheduler,
    *,
    horizon_s: float = 60.0,
    control_period_s: float = 1.0,
    expected_output_tokens: float = 128.0,
) -> OpenLoopPlan:
    """Optimize a nominal fluid schedule using a declared output-length forecast."""

    del scheduler  # The fluid plan assumes the fixed chunked scheduler named by the caller.
    plant.validate()
    if horizon_s <= 0.0 or control_period_s <= 0.0:
        raise ValueError("horizon and control period must be positive")
    horizon_steps = int(np.ceil(horizon_s / control_period_s))
    arrivals = _arrival_work_grid(
        workload,
        plant.profile,
        horizon_steps,
        control_period_s,
        expected_output_tokens,
    )
    profile = plant.profile
    minimum = profile.minimum_clock_mhz
    maximum = profile.maximum_clock_mhz
    service_levels = 0.5 * (
        profile.prefill_tokens_per_s / profile.prefill_tokens_per_s[-1]
        + profile.decode_tokens_per_s / profile.decode_tokens_per_s[-1]
    )
    minimum_service = float(service_levels[0])
    maximum_power = np.maximum(profile.prefill_power_w, profile.decode_power_w)
    normalized_power = maximum_power / maximum_power[-1]
    power_slope = (normalized_power[-1] - normalized_power[0]) / (
        1.0 - minimum_service
    )

    # Variables are n normalized service levels followed by n end-of-period
    # backlogs.  The inequalities implement b_{k+1} >= b_k + w_k - h u_k.
    number_variables = 2 * horizon_steps
    linear_cost = np.zeros(number_variables, dtype=float)
    linear_cost[:horizon_steps] = power_slope
    linear_cost[horizon_steps:] = 20.0
    linear_cost[-1] += 20.0
    constraint = np.zeros((horizon_steps, number_variables), dtype=float)
    bound = -arrivals.copy()
    for index in range(horizon_steps):
        constraint[index, index] = -control_period_s
        constraint[index, horizon_steps + index] = -1.0
        if index > 0:
            constraint[index, horizon_steps + index - 1] = 1.0
    solution = linprog(
        linear_cost,
        A_ub=constraint,
        b_ub=bound,
        bounds=[(minimum_service, 1.0)] * horizon_steps
        + [(0.0, None)] * horizon_steps,
        method="highs",
    )
    if solution.success:
        optimized_service = np.asarray(solution.x[:horizon_steps], dtype=float)
    else:
        optimized_service = np.ones(horizon_steps, dtype=float)
    continuous = np.interp(optimized_service, service_levels, profile.clock_mhz)
    applied = np.array(
        [plant.profile.quantize_clock(value, downward=True) for value in continuous]
    )
    predicted_backlog, _, _ = _fluid_rollout(
        applied,
        arrivals,
        plant,
        control_period_s,
    )
    from inference_serving import workload_checksum

    objective_value = float(solution.fun) if solution.success else float("nan")
    return OpenLoopPlan(
        time_s=np.arange(horizon_steps, dtype=float) * control_period_s,
        continuous_clock_mhz=continuous,
        applied_clock_mhz=applied,
        predicted_backlog_s=predicted_backlog,
        objective=objective_value,
        optimization_method="HiGHS linear program",
        success=bool(solution.success and np.all(np.isfinite(continuous))),
        message=str(solution.message),
        control_period_s=control_period_s,
        workload_checksum=workload_checksum(workload),
        profile_status=plant.profile.profile_status,
    )


class OpenLoopClockController:
    """Clock controller that replays a precomputed plan without feedback."""

    def __init__(self, plan: OpenLoopPlan):
        self.plan = plan
        self.plan_dt_s = plan.control_period_s
        self.plans_by_step = {0: tuple(float(value) for value in plan.applied_clock_mhz)}
        self.plan_start_times_by_step = {0: 0.0}
        self.__name__ = "open_loop_clock"

    def __call__(self, observation: ServingObservation) -> float:
        index = min(
            int(
                np.floor(
                    (observation.time_s + 1e-12)
                    / self.plan.control_period_s
                )
            ),
            self.plan.applied_clock_mhz.size - 1,
        )
        return float(self.plan.applied_clock_mhz[index])

Example: Gradient Descent with Momentum as DOCP¶

To connect this lens to familiar practice, including hyperparameter optimization, treat the learning rate and momentum (or their schedules) as controls. Rather than fixing them a priori, we can optimize them as part of a trajectory optimization. The optimizer itself becomes the dynamical system whose execution we shape to minimize final loss.

Program: gradient descent with momentum on a quadratic loss. We fit $\boldsymbol{\theta}\in\mathbb{R}^p$ to data $(\mathbf{A},\mathbf{b})$ by minimizing

The program maintains parameters $\boldsymbol{\theta}_k$ and momentum $\mathbf{m}_k$. Each iteration does:

1. compute gradient $\mathbf{g}_k=\nabla_{\boldsymbol{\theta}}\ell(\boldsymbol{\theta}_k)=\mathbf{A}^\top(\mathbf{A}\boldsymbol{\theta}_k-\mathbf{b})$

2. update momentum $\mathbf{m}_{k+1}=\beta_k \, \mathbf{m}_k + \mathbf{g}_k$

3. update parameters $\boldsymbol{\theta}_{k+1}=\boldsymbol{\theta}_k - \alpha_k \, \mathbf{m}_{k+1}$

State, control, and transition. Define the state $\mathbf{x}_k=\begin{bmatrix}\boldsymbol{\theta}_k\\ \mathbf{m}_k\end{bmatrix}\in\mathbb{R}^{2p}$ and the control $\mathbf{u}_k=\begin{bmatrix}\alpha_k\\ \beta_k\end{bmatrix}$. One program step is

Executing the program for $T$ iterations gives the trajectory

Objective as a DOCP. Choose terminal cost $c_T(\mathbf{x}_T)=\ell(\boldsymbol{\theta}_T)$ and (optionally) stage costs $c_k(\mathbf{x}_k,\mathbf{u}_k)=\rho_\alpha \, \alpha_k^2+\rho_\beta\,(\beta_k- \bar\beta)^2$. The program-as-control problem is

Backpropagation = reverse-time costate recursion. Because $\Phi_k$ is differentiable, reverse-mode AD computes $\nabla_{\mathbf{u}_{0:T-1}} \big(c_T+\sum c_k\big)$ by propagating a costate $\boldsymbol{\lambda}_k=\partial \mathcal{J}/\partial \mathbf{x}_k$ backward:

and the gradients with respect to controls are

Unrolling a tiny horizon ($T=3$) to see the composition:

What if the program branches? Suppose we insert a “skip-small-gradients” branch

which is non-differentiable because of the indicator. The DOCP view still applies, but gradients are unreliable. Two practical paths: smooth the branch (e.g., replace $\mathbf{1}\{\cdot\}$ with $\sigma((\lVert\mathbf{g}_k\rVert-\tau)/\epsilon)$ for small $\epsilon$) and use autodiff; or go derivative-free on $\{\alpha_k,\beta_k,\tau\}$ (e.g., SPSA or CMA-ES) while keeping the inner dynamics exact.

Variants: Lagrange and Mayer Problems¶

How does moving cost between running and terminal terms change the representation without changing the underlying control problem?

The Bolza form is general enough to cover most situations, but two common special cases deserve mention:

• Lagrange problem (no terminal cost) If the objective only accumulates stage costs:

Example: Energy minimization for a delivery drone. The concern is total battery use, regardless of the final position.

• Mayer problem (terminal cost only) If the objective depends only on the final state:

Example: Satellite orbital transfer. The only goal is to reach a specified orbit, no matter the fuel spent along the way.

These distinctions matter when deriving optimality conditions, but conceptually they fit in the same framework: the system evolves over time, and we choose controls to shape the trajectory.

Reducing to Mayer Form by State Augmentation¶

Although Bolza, Lagrange, and Mayer problems look different, they are equivalent in expressive power. Any problem with running costs can be rewritten as a Mayer problem (one whose objective depends only on the final state) through a simple trick: augment the state with a running sum of costs.

The idea is straightforward. Introduce a new variable, $y_t$, that keeps track of the cumulative cost so far. At each step, we update this running sum along with the system state:

where $\tilde{\mathbf{x}}_t = (\mathbf{x}_t, y_t)$. The terminal cost then becomes:

The overall effect is that the explicit sum $\sum_{t=1}^{T-1} c_t(\mathbf{x}_t,\mathbf{u}_t)$ disappears from the objective and is captured implicitly by the augmented state. This lets us write every optimal control problem in Mayer form.

This reduction serves two purposes. First, it often simplifies mathematical derivations, as we will see later when deriving necessary conditions. Second, it can streamline algorithmic implementation: instead of writing separate code paths for Mayer, Lagrange, and Bolza problems, we can reduce everything to one canonical form. That said, this unified approach is not always best in practice. Specialized formulations can sometimes be more efficient computationally, especially when the running cost has simple structure.

The unifying theme is that a DOCP may look like a generic NLP on paper, but its structure matters. Ignoring that structure often leads to impractical solutions, whereas formulations that expose sparsity and respect temporal coupling allow modern solvers to scale effectively. In the following sections, we will examine how these choices play out in practice through single shooting, multiple shooting, and collocation methods, and why different formulations strike different trade-offs between robustness and computational effort.

Summary and Outlook¶

A finite horizon converts a controlled dynamical model into a nonlinear program over states and actions. The differential-drag and inference examples also delimit its guarantee: feasibility and performance apply to the stated initial condition, dynamics, and disturbance forecast. Bolza, Lagrange, and Mayer forms change the bookkeeping without changing the admissible action sequences.

Under a constraint qualification, the KKT conditions are necessary at a local solution of the resulting nonlinear program, but their raw form hides the direction of time. Can the multipliers be organized into a backward recursion that matches the forward state dynamics? Adjoints and the discrete-time Pontryagin principle provide that organization.

Exercises¶

Self-checks¶