Search code examples
pythondrake

MultibodyPlant’s dynamics as part of a Python function/lambda constraint in a MathematicalProgram for direct collocation from scratch

I'm trying to implement direct collocation from scratch in a MathematicalProgram such that each constraint is a Python function, meaning that I can get the gradient of each constraint and cost with respect to their inputs. My goal is to use these gradients for a downstream task.

I'm converting the "Direct Collocation for the Pendulum" part of this notebook: https://github.com/RussTedrake/underactuated/blob/master/trajopt.ipynb to a MathematicalProgram. I've been able to convert the cost and constraints to lambdas instead of Formulas. For instance:

torque_limit = 3.0  # N*m.
u = dircol.input()
dircol.AddConstraintToAllKnotPoints(-torque_limit <= u[0])
dircol.AddConstraintToAllKnotPoints(u[0] <= torque_limit)

becomes

torque_limit = 3.0  # N*m.
for n in range(N-1):
    constraint = prog.AddConstraint(lambda u: u, [-torque_limit], [torque_limit], [u[n][0]])

This means I can use ExtractGradient with that constraint:

ExtractGradient(constraint.evaluator().Eval(InitializeAutoDiff(arbitrary_u_vector)))

The one constraint I have not been able to turn into a Python lambda/function is the constraint for Pendulum's dynamics, since that is implemented inside the black box of DirectCollocation. I've read the C++ for DirectCollocation, but I haven't been able to figure it out.

What would it take to either reimplement trajectory optimization's dynamics constraint in Python as a function/lambda or access the dynamics constraint's gradient in Python in some other way? I would like to be able to do this for an arbitrary MultibodyPlant, say Quadrotor or Acrobot, not just Pendulum.

Solution

  • It should definitely be possible to implement the DirectCollocation constraint in python. This notebook has a bunch of relevant examples of using python functions as constraints; it could help. But it should be a fairly straightforward port of the DirectCollocationConstraint::Eval you've found in C++.

    But you can also always get the gradients from the constraints that are being added by DirectCollocation. That class is just a helper class that makes it easier to set up the MathematicalProgram for you. You can still get that program back out and evaluate any of the constraints individually if you like. I'm not convinced you get anything by reimplementing it yourself, and will certainly have worse performance in python.