[General] GPU Accelerated Trajectory Optimization with External Transcription

[andrewrosemberg]

@jbcaillau and @ocots,

I am trying to figure out the best way to implement optimal control problems for trajectory optimization in Julia to be solved with GPUs, making the best use of existing packages.

After chatting with @amontoison, I believe that OptimalControl.jl can be a useful framework to communicate with the appropriate modeling tools (e.g., Examodels) and solvers (e.g., those in MadSuit), besides handling the integrator part.

Unfortunately, I haven’t been able to completely derive the dynamics of my quadruped robot from scratch, which I believe is necessary for OptimalControl.jl.

Taking inspiration from TORA.jl (cc @ferrolho), I thought of using RigidBodyDynamics.jl (RBD) to derive them from a URDF file. In TORA.jl, this approach calls the RBD dynamics inference API compatible with AD to calculate both the Jacobian and Hessian. From my chat with @amontoison, I believe this could be made compatible with OptimalControl.jl by using the NLPModels Manual API: https://jso.dev/tutorials/create-a-manual-model/. The issue here is that this approach would require me to extend and ensure that RigidBodyDynamics.jl can perform dynamics inference on the GPU, as far as I know (AFAIK).

A while back, I also tried to rewrite RigidBodyDynamics.jl to directly create the appropriate constraints in a given JuMP model instead of solving the dynamics (andrewrosemberg/RigidBodyDynamics.jl#1). However, it seems that my naive implementation produces extremely large expressions, causing JuMP to stall when trying to check if they belong to the model. I now believe the rewrite would need to be much more intricate, ensuring that the expressions are created cleverly, avoiding redundancy, and being appropriately compatible with how the GPU kernels should be created for good performance.

Sorry for the long post, but I wanted to see what you and others from the community think. :)

[amontoison]

@ocots @jbcaillau bump

[jbcaillau]

hi @andrewrosemberg sorry for the late reply. busy these weeks (including next one 🥲). not sure to completely grasp your issue, but i guess i have a global picture of it. could you please:

• tell us state and control dimensions (how big is it)?
• give us an idea of the dynamics? is it an implicit ODE that you generate using AD, e.g. because you differentiate some mechanical constraints...?

my impression is that things should be split, without giving up about performance anywhere:

• math modelling (e.g. using a high level enough DSL such as the one from OptimalControl.jl) (*)
• discretisation (direct transcription into an NLP) using a suitable modeller / framework
• running (optimising) on CPU / GPU efficiently (with subexpressions fartorisation, leveraging sparsity...)

(*) currently not accepting implicit dynamics / DAE, though

[andrewrosemberg]

Thank you @jbcaillau ,

tell us state and control dimensions (how big is it)?

My question is for general problems, but one problem that I have been working on is the Atlas robot (state size of 58 and action around 27). The idea is to scale for multi-agent, so this would increase significantly.

give us an idea of the dynamics? is it an implicit ODE that you generate using AD, e.g. because you differentiate some mechanical constraints...?

Currently, I am using RK4 as the integrator, but I am easy to change if there is a more efficient/robust one.

---

The issue that I would like to address is that $\dot x = f(x, u)$ is computed by an external call (RigidbodyDynamics.jl - that also does the AD for the dynamics), but OptimalControl.jl has no API for external calls yet, from what I understand. Do you see a path for implementing this?

[jbcaillau]

hi @andrewrosemberg actually OC.jl does have external calls, with the default backend that is with ADNLPModels. check the example below where $dot{x} = f(x,u)$ with $f$ an arbitrary Julia function. just needs to be differentiable with some AD backend such as ForwardDiff.jl.

https://control-toolbox.org/Kepler.jl/stable/#Direct-solve

[amontoison]

@jbcaillau Can we plug in any AbstractNLPModel?
If it is the case, it will help a lot @andrewrosemberg.

Update: If I am not wrong, it is you that create the ADNLPModel 🤔

[jbcaillau]

@amontoison yes, we create the (concrete) ADNLModel in CTDirect (@PierreMartinon)

@jbcaillau Can we plug in any AbstractNLPModel? If it is the case, it will help a lot @andrewrosemberg.

Update: If I am not wrong, it is you that create the ADNLPModel 🤔