Duplication of objects passed to OpenMDAO ExplicitComponent in Dymos phase?

[FluffyCodeMonster]

This is related to my previous question: #1138, but I think I've been able to isolate the potential issue after running more tests.

• When calculating the dynamics (state rate) values, my code incorporates information about the vehicle's environment, which is a function of the state at the current time.
• The environment values can be defined analytically, as a simple function of the state, or they can come from an interpolator object, which interpolates the environment values at the state from real-world input. The interpolator objects are generated offline and then pickled. Each one is about 1.5GB in size.
• When the program is run, it has to pass the environment data to the OpenMDAO ExplicitComponent which defines the dynamics model, so that it can be used in the dynamics calculations.
• When I run the optimisation with the lightweight, analytic environment model (see example code below), it runs absolutely fine, uses hardly any memory and completes a bout of grid refinement without issue.
• When the optimisation is run with an interpolated environment, the process (Dymos running in Windows Subsystem for Linux) saturates 280GB of RAM and then crashes. (With sufficiently few collocation nodes, it successfully completes an optimisation, but crashes on the grid refinement)
• I thought the issue could have been related to the colouring or Jax compilation operations, but the memory is saturated and the process is killed even after disabling these.
• The colouring with the analytic environment model should also surely be the same as with the interpolated model (all of the dynamics should be excited in both cases), so I don't think it can be the colouring.

This leads me to the questions:

• Is Dymos creating duplicates of the environment model that I'm passing in? E.g. one for each collocation point? This would explain why the memory usage is so high.
• Am I passing the environment model to the OpenMDAO ExplicitComponent in the correct way?
• If Dymos is creating duplicates, is there any way that I can stop it?

In the actual code below, I don't pass the Environment object itself but the reference to the 'get_vals' method. I thought that this might help if there is copying under the hood, e.g. only the method reference would be copied and not the whole object, but unfortunately the memory is still used up and the process is killed.

A simplified version of the code below shows the example dynamics, environment code and problem setup. Please let me know if anything's unclear.

Thank you for any assistance.

Environment code:

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp

class Environment:
  def __init__(self):
    # Initialisation code

  def get_vals(self, state_vec):
    pass

class ExampleAnalyticEnvironment(Environment):
  def get_vals(self, state_vec):
    a = ...
    b = ...
    c = ...
    return jnp.array([a, b, c])

class InterpolatedEnvironment(Environment):
  def __init__(self, interpolator_path):
    # Interpolator data loaded from file
    # Interpolator object created  (wrapper for Interpax interpolators)
    self.interp = ...

  def get_vals(self, state_vec):
    return self.interp.interpolate(state_vec)

Dymos dynamics model code:

# For dynamics code
import OtherClass

import openmdao.api as om
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp

# Dynamics model
class Dynamics(om.ExplicitComponent):
  def __init__(self, **kwargs):
    super().__init__(**kwargs)

    # Jax jit _compute_primal for faster performance
    self._compute_primal_vec = jax.vmap(self._compute_primal)
    # Autodiff for automatic computation of partial derivatives
    self._compute_partials_vec = jax.jit(jax.vmap(jax.jacfwd(self._compute_primal, argnums=np.arange(3))))

  def initialize(self):
    self.options.declare('num_nodes', types=int)
    # ************** Function to compute environment values at given state passed to model **************
    self.options.declare('get_env_vals')

  def setup(self):
    nn = self.options['num_nodes']

    # States
    self.add_input('example_state_1', shape=(nn,)) # units also provided in real problem code
    self.add_input('example_state_2', shape=(nn,)) # units also provided in real problem code

    # Controls
    self.add_input('example_control', shape=(nn,)) # units also provided in real problem code

    # Output state rates
    self.add_output('example_state_1_dot', val=jnp.zeros(nn)) # units also provided in real problem code
    self.add_output('example_state_2_dot', val=jnp.zeros(nn)) # units also provided in real problem code

    # Partials declared analytically
    arange = jnp.arange(nn)
    self.declare_partials(of='*', wrt='*', method='exact', rows=arange, cols=arange)
    
    # Dynamics go here
    @partial(jax.jit, static_argnums=(0,))
    def _compute_primal(self, example_state_1, example_state_2, example_control):
      # Dynamics defined in static method of another class
      ************** Pass get_env_vals function **************
      grads = OtherClass.compute_gradients(self.options['get_env_vals'], example_state_1, example_state_2, example_control)
      return grads

  def compute(self, inputs, outputs):
    example_state_1_dot, example_state_2_dot = jnp.array(self._compute_primal_vec(*inputs.values()))
    outputs['example_state_1_dot'] = example_state_1_dot
    outputs['example_state_2_dot'] = example_state_2_dot

  def compute_partials(self, inputs, partials):
    output_names = ['example_state_1_dot', 'example_state_2_dot']
    input_names = ['example_state_1', 'example_state_2', 'example_control']

    # Compute analytic partial derivative values using Jax
    computed_partials = self._compute_partials_vec(*inputs.values())

    # Cycle through computed partials
    for out_ind, output_name in enumerate(output_names):
      for in_ind, input_name in enumerate(input_names):
        partials[output_name, input_name] = computed_partials[out_ind][in_ind]

Setting up the Dymos model to optimise:

import openmdao.api as om
import dymos as dm

# Create environment object
env = ExampleAnalyticEnvironment()
# or..
env = InterpolatedEnvironment("path/to/interpolator/file")

prob = om.Problem(reports=None)

traj = dm.Trajectory()

# Trajectory gets added to the problem
prob.model.add_subsystem('traj', traj)

# Define transcription. num_segments and order defined elsewhere.
tx = dm.GaussLobatto(num_segments=num_segments, order=order, compressed=False)

# Create phase
************** Pass Environment.get_vals function **************
phase = dm.Phase(ode_class=Dynamics, transcription=tx, ode_init_kwargs={'get_env_vals': env.get_vals})

traj.add_phase(dynamics_name, dynamics_phase)

# ... Define states, controls, objective, optimiser settings, declare_colouring(), etc ..

prob.setup()

# ... Set initial values ...

dm.run_problem(prob, run_driver=True, refine_iteration_limit=1, refine_method='hp')

[robfalck]

We're still looking into this. ExplicitShooting should be instantiating two copies of the ODE...one is used in the wrapper for scipy,integrate, and the other is a vectorized one that is used to compute any ODE outputs at all of the nodes once the states have been propagated. We're still trying to track down the source of the memory issues here.

Also, I noticed you're using jax. By far the easiest way to do this with OpenMDAO as of December 2024 is to use JaxExplicitComponent and JaxImplicitComponent. There's an example of a brachistochrone implementation that utilizes it here: https://openmdao.github.io/dymos/examples/implicit_duration/brachistochrone_implicit_duration.html#solving-the-brachistochrone-without-an-optimizer

There are a few hoops to jump through if you need static component attributes or options in your compute_primal function that are explained here: https://openmdao.org/newdocs/versions/latest/features/experimental/jax_explicitcomp_api.html

[FluffyCodeMonster]

Thank you for investigating the issue, and for the links to the JaxExplicitComponent code - I didn't know about this, and it looks to be doing something very similar to my code above (is there anything obviously wrong with my approach?). I'd like to try to avoid significant code revisions if I can (also it looks like it uses a newer version of Dymos than I'm on and I don't want to risk breaking things at this point...), but I'll definitely investigate it if I can't find a solution with the existing code. Your answer raised an interesting question about the environment interpolator being interpreted as static though (which it should be, because it doesn't change).

In my code, there are two objects passed to the OpenMDAO ExplicitComponent - a vehicle parameters object, and an environment interpolator object. I've checked and both are hashable (a requirement of static parameters), and the hash value doesn't change after calling the interpolator, so there don't seem to be any internal state changes which is good. I've modified the code so that both are now passed to _compute_primal(...) as arguments,

@partial(jax.jit, static_argnums=(0, 1, 2))
def _compute_primal(self, vehicle_model, environment_interp, example_state_1, example_state_2, example_control):
  ...

and the vectorisation and auto-diff code in __init__ has become

# Nones for the vehicle model and environment interpolator
in_axes = [None, None, 0, 0, 0]
self._compute_primal_vec = jax.vmap(self._compute_primal, in_axes=in_axes)
# 'argnums=np.arange(2, 5)' - don't try to differentiate the vehicle model or environment interpolator
self._compute_partials_vec = jax.jit(jax.vmap(jax.jacfwd(self._compute_primal, argnums=np.arange(2, 5)), in_axes=in_axes), static_argnums=(0, 1))

I was optimistic that this would fix it, but unfortunately it didn't really affect the performance and the issues still persisted.

Extra troubleshooting information 1)
This is not the first time I've faced the program being killed due to memory issues. My initial code looped through a set of different environment cases (simpler optimisations, without linked phases, which were completing successfully with a grid refinement) in a Python loop, but this only ever made it through one optimisation before the process was killed. (My solution was to loop in bash instead, calling the Python script with different command line arguments each time)

Extra troubleshooting information 2)
There is always quite a delay between the program being started and the first 'output line' of the optimisation. I always attributed this to the colouring (which can take a couple of minutes) and the Jax compilation, although this initialisation process can still be slow without it. I did receive this error earlier:

Although my computer was also running CFD so it might have been due to the other present computational load.

Extra troubleshooting information 3)
I'm running the following software versions:

• dymos 1.11.0
• openmdao 3.34.2
• jax 0.4.31
• jaxlib 0.4.31

Sorry for the bombardment of information, I'm just trying to include any details which might be relevant in diagnosing the problem. Many thanks again for your assistance.

[FluffyCodeMonster]

Hi, I was just wondering if there's any new insight on this? Many thanks.

[kaushikponnapalli]

Have you tried using om.MetaModelStructuredComp to do your interpolation? You would have to turn your ODE into a group where the meta model computes your environment values and an explicit component computes your dynamics. A basic version of this approach is in the min time climb example. In the propulsion group we interpolate mach number and altitude to determine the maximum thrust

[FluffyCodeMonster]

Hi, thank you for the suggestion, that sounds like a good idea. Would this still work with Jax (e.g. for automatic differentiation of the interpolated quantities?). My code mostly works and I'm trying to avoid significant revisions at this point in time if I can help it, but I'll definitely look into it as a solution.