Dymos performance question

[FluffyCodeMonster]

I'm experiencing some potential optimisation issues using Dymos (with IPOPT as the solver), and wanted to ask whether anything seems wrong with my setup. I'm concerned as my optimisation is taking a very long time compared to the Dymos examples, and I don't know if I'm doing something wrong which might be slowing it unnecessarily. I've tried to compare run times against systems of a similar size/complexity in other papers and it doesn't seem too far wrong (assuming they're set up well), but thought it would be a good idea to seek insight here.

The system I'm optimising has 12 states (plus other OpenMDAO states for output), with quite a lot of coupling in the differential equations, and two control inputs. The partial derivates are calculated using Jax, and each iteration of the dynamics involves an interpolation, performed by the Interpax Python library. I've checked the partials with the check_partials method in Dymos and the gradient errors look okay (the largest absolute error is 0.00012, much smaller for complex step, although this only works if I disable the interpolation so it's not a true representation). With IPOPT it can take hours to solve (usually with one grid refinement), and even still often ends with an error message, or only 'solves to an acceptable level'. The outputs of some of my runs are shown in the attached spreadsheet - as can be seen, mostly all of them exit with an IPOPT error.

The primal error often gets down to e-13/14 or so but the dual error struggles, getting stuck at about e-3/-4 at best, and animating the state discretisation nodes/collocation points shows that they do find a good solution quite quickly but then stay put and jitter for the rest of the time. However if the gradients are okay I'm not sure what's happening. This might just be the normal behaviour of a system like this, but I don't really have a reference point for these kinds of optimisation problems.

For extra information, the Jacobian shape pre-grid-refinement is (6374, 5249) - it takes ~99s to compute the sparsity and ~18s to compute the colouring. Post-refinement it's (9824, 6449); it takes ~156s to compute the sparsity and ~46s for the colouring. All of this is being done on a very powerful desktop computer (although the software as written doesn't utilise many of the cores - this is another peculiarity, because on my personal desktop it uses all 16 of them, but it's perhaps to do with the Jax/XLA compatibility of the machine).

I was wondering if these solve times (also shown in the spreadsheet) look normal, and if it's standard to encounter errors like this, or whether I'm doing something fundamentally wrong? Many thanks for any assistance - I really appreciate the help.

Investigating solver performance.xlsx

[FluffyCodeMonster]

If it helps, I can also provide the iteration-by-iteration output from a solve in IPOPT.

[robfalck]

Even with good gradients, its possible that there are multiple paths that provide the same (or nearly the same) performance at that point.

Nonlinear optimizers like SNOPT, IPOPT, and SLSQP typically struggle at these conjugate points because they're tuned to find a minimum that might not be well defined.

Some things you can try:

• You can just give up the tolerance and let it converge the optimality loosely.
• You can add some penalty term to the objective that provides some extra curvature to it so that it prefers some solution. For instance, if finding a minimum fuel burn you might instead define the objective as fuel_burn + 0.1 * time.
• Insensitivity to the controls can also do this. One resolution is to use the lower-order polynomial controls in dymos that simply don't give that much freedom to the optimizer. Alternatively, you can formulate an arclength of the control polynomial in time and add a penalty term to that in the objective so that it prefers a certain result in the absense of any other information.

[FluffyCodeMonster]

Hi Rob, thank you for your help and insight - it's always very valuable - and apologies for my delayed response. If I could clarify a couple of things:-

For the first point, do you mean to relax the dual error, e.g. let the optimisation stop when the dual is at e.g. 10^-3? I suppose this would make it hard to make claims about the optimality of the solutions, because there is still scope in the problem for the objective function to be much lower?

For the second and third points, I do actually have some additional curvature in the objective function, as it's trying to minimise inertial energy at the endpoint plus the squared integral of the controls along the trajectory. Really I want to minimise the rate of change of the controls along the trajectory rather than their squared integral (this is what I was trying to put in the SO question which led to this dicussion, although not expressed very clearly), so that I get nice smooth control traces, as they do have a tendancy to jitter, although your last point seems like an interesting resolution to this and something I should try. Do you think all of this control flexibility might be what's causing the optimiser to find a rough solution and then waste a lot of time 'jittering', with high dual error, just because it has so much scope to vary the controls? It does sound likely now you mention it... I always assumed that the order of the polynomials was the spline order between the state discretisation nodes, but I think it works differently for the Gauss-Lobatto collocation used in Dymos (since it's orthogonal collocation) - I'll have to read up on it in the paper linked in the Dymos logs.

[FluffyCodeMonster]

I found the low-order polynomial control option, as in this example: https://openmdao.github.io/dymos/examples/ssto_moon_polynomial_controls/ssto_moon_polynomial_controls.html, and have tried it. I still get some strange behaviour but I'll experiment more. I suppose another valid approach would be to have the controls on a separate transcription connected to the system dynamics, so that the dynamics are solved with high fidelity without the control space being too large (e.g. maybe 1 control node every 5-10 dynamics nodes)? I'll have a think about how to accomplish this.

[robfalck]

I always assumed that the order of the polynomials was the spline order between the state discretisation nodes, but I think it works differently for the Gauss-Lobatto collocation used in Dymos (since it's orthogonal collocation) - I'll have to read up on it in the paper linked in the Dymos logs.

For each segment in Gauss-Lobatto, the controls are defined at all nodes. For a 3rd order segment with two endpoint nodes and a single midpoint collocation node, this results in a quadratic control fit.

For consistency, we do the same thing in Radau. The controls are defined at the first N-1 nodes in the segment with the value at the endpoint being the fit of the first 3 extrapolated to the endpoint.

Are you enforcing control rate continuity between segments? Doing so could be adding some Gibbs phenomenon to the solution and causing some of your ringing.

As for having the controls at a lower fidelity, that's effectively what the polynomial control is doing, but you're right in that there are a lot of different ways that could be parameterized.

[FluffyCodeMonster]

Hi Rob, thank you, I tried using the polynomial controls and it worked well, although I had to use quite high orders to give it the required flexibility (about 25th order) as it's optimising quite a long trajectory, and in that case it still exited early with a 'EXIT: Solved To Acceptable Level. Cannot call restoration phase at almost feasible point, but acceptable point from iteration 1818 could be restored.' and a relatively high dual error of 8.54E-04. The runs I tried with the lower order polynomials though, such as order 7, finished quickly and with dual errors of E-09, which shows that it works (thank you!).

I've also tried putting the controls on a tandem phase, which also works. My idea to tackle the ringing was to penalise the squared integral of the control rate in the objective (similar to what's discussed above), hence I tried using control rate as the input. This worked well - here it is for the double integrator problem (note the difference in the transcriptions of the controls and the dynamics):

Here the integral of u_dot^2 (the label in the plot is missing the square) is a state which can just be penalised in the objective at the final time point.

Sadly doing rate-based control really doesn't seem to scale up well to my real system/trajectory optimisation, so it looks like I'm back to specifying the control values directly (example for the double integrator):

I'd still like to be able to compute the integral of u_dot^2 to penalise the control jitter, but no longer have access to u_dot. This is the situation I was trying to describe in my original SO question. I don't know if there's a way that I can differentiate u (the control) in Dymos/OpenMDAO (the ExplicitComponent which contains the model)? If I had access to the times of the state discretisation and collocation nodes I could maybe use something like the Numpy gradients function and the Scipy/OpenMDAO integrators. I've looked at this question: https://stackoverflow.com/questions/77698377/dymos-how-to-specify-the-values-of-a-time-varying-variable, but am not quite sure whether it's the right thing.

Also yes, I do have control rate continuity turned on. I'll try turning it off to see if it helps.

I apologise for the string of questions - I feel like this is very nearly there. Many thanks again for all of your help. Best, Freddie.

[FluffyCodeMonster]

I made an initial attempt at calculating the control rates for the simple double integrator problem shown above, which is the following:

import matplotlib.pyplot as plt
import openmdao.api as om
import dymos as dm
from dymos.examples.plotting import plot_results

import numpy as np
from pathlib import Path

# Code from Dymos example: https://openmdao.github.io/dymos/examples/double_integrator/double_integrator.html
class DynamicsODE(om.ExplicitComponent):
    """
    The double integrator is a special case where the state rates are all set to other states
    or parameters.  Since we aren't computing any other outputs, the ODE doesn't actually
    need to compute anything.  OpenMDAO will warn us that the component has no outputs, but
    Dymos will solve the problem just fine.

    Note we still have to declare the num_nodes option in initialize so that Dymos can instantiate
    the ODE.

    Also note that neither time, states, nor parameters have targets, since there are no inputs
    in the ODE system.
    """

    def initialize(self):
        self.options.declare('num_nodes', types=int)
    
    def setup(self):
        nn = self.options['num_nodes']

        self.add_input('u_dot', shape=(nn,), units='m/s**3')

        self.add_output('u_', val=np.zeros(nn), units='m/s**2')

class ControlODE(om.ExplicitComponent):
    def initialize(self):
        self.options.declare('num_nodes', types=int)
    
    def setup(self):
        nn = self.options['num_nodes']
    
        self.add_input('u', shape=(nn,), desc='u', units='N/kg')

        self.add_output('u_int_rate', val=np.zeros(nn), units='N/kg')

        self.declare_partials('u_int_rate', 'u', method='cs')

    def compute(self, inputs, outputs):
        u = inputs['u']

        # Square control rates
        u_sq = np.square(u)

        outputs['u_int_rate'] = u_sq


p = om.Problem(model=om.Group())

p.driver = om.pyOptSparseDriver()
p.driver.options['optimizer'] = 'IPOPT'
p.driver.opt_settings['print_level'] = 5
p.driver.declare_coloring()

# Setup the trajectory and its phase
traj = p.model.add_subsystem('traj', dm.Trajectory())

# The transcription of the first phase
control_tx = dm.GaussLobatto(num_segments=5, order=3, compressed=False)

# The transcription for the second phase (and the secondary timeseries outputs from the first phase)
dynamics_tx = dm.GaussLobatto(num_segments=20, order=3, compressed=False)

# First phase - controls, low-speed

control_phase = dm.Phase(ode_class=ControlODE, transcription=control_tx)
traj.add_phase('control_phase', control_phase)

control_phase.set_time_options(fix_initial=True, duration_bounds=(.5, 10))

control_phase.add_control('u', shape=(1,), continuity=True, rate_continuity=True, units='N/kg', lower=-1., upper=1.)
control_phase.add_state('u_int', fix_initial=True, rate_source='u_int_rate', units='N*s/kg')  # m/s

# Add alternative timeseries output to provide control inputs for the next phase
control_phase.add_timeseries('control_to_dyn_ts', transcription=dynamics_tx, subset='control_input')

# Second phase - dynamics

dynamics_phase = dm.Phase(ode_class=DynamicsODE, transcription=dynamics_tx)
traj.add_phase('dynamics_phase', dynamics_phase)

dynamics_phase.set_time_options(fix_initial=True, input_duration=True)

dynamics_phase.add_state('v', fix_initial=True, fix_final=True, rate_source='u', units='m/s')
dynamics_phase.add_state('d', fix_initial=True, fix_final=False, rate_source='v', units='m')

dynamics_phase.add_control('u', shape=(1,), opt=False, units='m/s**2')

# ===
# Attempting to calculate control rate
dynamics_phase.add_control('u_dot', units='m/s**3')
dynamics_phase.add_state('u_', rate_source='u_dot', units='m/s**2')
# Ensure that u and u_ are equal.
dynamics_phase.add_path_constraint('ud = u - u_', equals=0)
# ===

# Connect the two phases
p.model.connect('traj.control_phase.t_duration_val', 'traj.dynamics_phase.t_duration')
p.model.connect('traj.control_phase.control_to_dyn_ts.u', 'traj.dynamics_phase.controls:u')

# Minimize time
dynamics_phase.add_objective('d', loc='final', scaler=-1)

p.setup()

control_phase.set_time_val(initial=0.0, duration=2.0)
control_phase.set_control_val('u', [0, 0])

dynamics_phase.set_state_val('d', [0, 1])
dynamics_phase.set_state_val('v', [0, 0])

p['traj.control_phase.states:u_int'] = 0.0

res = dm.run_problem(p, simulate=True, simulation_record_file='./dymos_simulation.db')
print(res)

sim = traj.simulate()

def plot(ax, phase_name, var_name, label):
    sol_times = p.get_val(f'traj.{phase_name}.timeseries.time')
    sim_times = sim.get_val(f'traj.{phase_name}.timeseries.time')
    sol_vals = p.get_val(f'traj.{phase_name}.timeseries.{var_name}')
    sim_vals = sim.get_val(f'traj.{phase_name}.timeseries.{var_name}')

    ax.scatter(sol_times, sol_vals, c='orange')
    ax.plot(sim_times, sim_vals)
    ax.set_xlabel('time (s)')
    ax.set_ylabel(label)

fig, axs = plt.subplots(5)

plot(axs[0], 'dynamics_phase', 'd', 'd $(m)$')
plot(axs[1], 'dynamics_phase', 'v', 'v $(m/s)$')
plot(axs[2], 'control_phase', 'u', 'u $(m/s^2)$')
plot(axs[3], 'control_phase', 'u_int', r'$ \int u$ $(m/s)$')
plot(axs[4], 'dynamics_phase', 'u_dot', r'$\dot{u}$ $(m/s^2)$')

fig.suptitle("Solution")

plt.show(block=True)

However, I get that there are 'not enough degrees of freedom' in the formulation. I would expect the optimiser to be adjusting the values of the u_dot control to satisfy the dynamics_phase.add_path_constraint('ud = u - u_', equals=0) constraint, but there must be something wrong with this approach. Alternatively, I could try putting the difference between u and the calculated u_ into the objective function. I did attempt this, but realised that it was only comparing u and u_ at the final point, so I'll need to try again. Do you think that either of these approaches are along the right lines for calculating the control rates at the state discretisation nodes and collocation points, or is this entirely the wrong angle?