Approach for Linearizing Nonlinear System around Decision Variables (x, u) from MathematicalProgram at each discrete point in time

Question

This is a follow up question on the same system from the following post (for additional context). I describe a nonlinear system as a LeafSystem_[T] (using templates) with two input ports and one output port. Then, I essentially would like to perform direct transcription using MathematicalProgram with an additional cost function that is dependent on the linearized dynamics at each time step (and hence linearized around the decision variables). I use two input ports as it seemed the most straightforward way for obtaining the linearized dynamics of the form from this paper on DIRTREL (if I can take the Jacobian with respect to input ports)

δx_i+1 ≈ A_iδx + B_iδu + G_iw

where i is the timestep, x is the state, the first input port can encapsulate u, and the second input port can model w which may be disturbance, uncertainty, etc.

My main question is what would be the most suitable set of methods to obtain the linearized dynamics around the decision variables at each time step using automatic differentation? I was recommended trying automatic differentiation after attempting a symbolic approach in the previous post, but am not familiar with the setup for doing so. I have experimented with

1. using primitives.Linearize() (calling it twice, once for each input port) which feels rather clunky and I am not sure whether it is possible to pass in decision variables into context

2. perhaps converting my system into a multibody and making use of multibody.tree.JacobianWrtVariable()

3. or formatting my system dynamics so that I can pass them in as the function argument for forwarddiff.jacobian

but have met limited success.

Answer 1

Sorry for my belated reply.

Since your dynamics can be written by hand, then I would suggest to create a templated function to compute Ai, Bi, Gi as

template <typename T>
void ComputeLinearizedDynamics(
 const LeafSystem<T>& my_system,
 const Eigen::Ref<const drake::VectorX<T>>& x,
 const Eigen::Ref<const drake::VectorX<T>>& u,
 drake::MatrixX<T>* Ai,
 drake::MatrixX<T>* Bi,
 drake::MatrixX<T>* Gi) const;

You will need to write down the matrix Ai, Bi, Gi by hand within this function. Then when you instantiate your LeafSystem with T=AutoDiffXd, this function will compute Ai, Bi, Gi with its gradient, given the state x, input u and disturbance w.

Then in the cost function, you could consider to create a sub-class of Cost class as

class MyCost {
 public:
  MyCost(const LeafSystem<AutoDiffXd>& my_system) : my_system_{&my_system} {}

 protected:
  void DoEval(const Eigen::Ref<const Eigen::VectorXd>& x_input, Eigen::VectorXd* y) const {
  // The computation here is inefficient, as we need to cast
  // x_input to Eigen vector of AutoDiffXd, and then call
  // DoEval with AutoDiffXd version, and then convert the 
  // result back to double. But it is easy to implement.
    const AutoDiffVecXd x_autodiff = math::initializeAutoDiff(x_input);
    AutoDiffVecXd y_autodiff;
    this->DoEval(x_autodiff, &y_autodiff);
    *y = math::autodiffToValueMatrix(y_autodiff);
  }

  void DoEval(const Eigen::Ref<const drake::AutoDiffVecXd>& x_input, drake::AutoDiffVecXd* y) const {
    // x_input here contains all the state and control sequence The authors need to first partition x_input into x, u
    drake::VectorX<T> x_all = x_input.head(num_x_ * nT_);
    drake::VectorX<T> u_all = x_input.tail(num_u_ * nT_);
    y->resize(1);
    y(0) = 0;
    // I assume S_final_ is stored in this class.
    Eigen::MatrixXd S = S_final_;
    for (int i = nT-1; i >= 0; --i) {
      drake::MatrixX<AutoDiffXd> Ai, Bi, Gi;
   
      ComputeLinearizedDynamics(
        *my_system_,
         x_all.segment(num_x_ * i, num_x_),
         u_all.segment(num_u_ * i, num_u_),
         &Ai, &B_i, &Gi);
      S = Ai.T*S + S*Ai + ... // This is the Riccati equation.
      // Now compute your cost with this S
      ...
    }
  }

  void DoEval(const Eigen::Ref<const VectorX<symbolic::Variable>& x, VectorX<symbolic::Expression>* y) const {
    // You don't need the symbolic version of the cost in nonlinear optimization.
    throw std::runtime_error("Not implemented yet");
  }
 private:
  LeafSystem<AutoDiffXd>* my_system_;
};

The DoEval function with autodiff version will compute the gradient of the cost for you automatically. Then you will need to call AddCost function in MathematicalProgram to add this cost together with all of x, u as the associated variable of this cost.

Answer 2

The easiest way to get Ai, Bi is to instantiate your system with AutoDiffXd, namely LeafSystem<AutoDiffXd>. The following code will give you Ai, Bi

MyLeafSystem<AutoDiffXd> my_system;
Eigen::VectorXd x_val = ...
Eigen::VectorXd u_val = ...
Eigen::VectorXd w_val = ...
// xuw_val concantenate x_val, u_val and w_val
Eigen::VectorXd xuw_val(x_val.rows() + u_val.rows() + w_val.rows());
xuw_val.head(x_val.rows()) = x_val;
xuw_val.segment(x_val.rows(), u_val.rows()) = u_val;
xuw_val.segment(w_val.rows()) = w_val;
// xuw_autodiff stores xuw_val in its value(), and an identity matrix in its gradient()
AutoDiffVecXd xuw_autodiff = math::initializeAutoDiff(xuw_val);

AutoDiffVecXd x_autodiff = xuw_autodiff.head(x_val.rows());
AutoDiffVecXd u_autodiff = xuw_autodiff.segment(x_val.rows(), u_val.rows());
AutoDiffVecXd w_autodiff = xuw_autodiff.tail(u_val.rows());

// I suppose you have a function x[n+1] = dynamics(system, x[n], u[n], w[n]). This dynamics function could be a wrapper of CalcUnrestrictedUpdate function.
AutoDiffVecXd x_next_autodiff = dynamics(my_system, x_autodiff, u_autodiff, w_autodiff);
Eigen::MatrixXd x_next_gradient = math::autoDiffToGradientMatrix(x_next_autodiff);
Eigen::MatrixXd Ai = x_next_gradient.block(0, 0, x_val.rows(), x_val.rows());
Eigen::MatrixXd Bi = x_next_gradient.block(0, x_val.rows(), x_val.rows(), u_val.rows());
Eigen::MatrixXd Gi = x_next_gradient.block(0, x_val.rows() + u_val.rows(), x_val.rows(), w_val.rows());

So you get the value of Ai, Bi, Gi in the end.

If you need to write a cost function, you will need to create a subclass of solvers::Cost. Inside the Eval function of this derived class, you will implement your code to first compute Ai, Bi, Gi, and then integrate the Riccati equation.

But I think since your cost function depends on Ai, Bi, Gi, the gradient of your cost function will depend on the gradient of Ai, Bi, Gi. Currently we don't provide the function to compute the second order gradient of the dynamics.

How complicated is your dynamical system? Is it possible to write down the dynamics by hand? If so, there are some shortcuts we can do to generate the second order gradient of your dynamics.

@sherm or other Drake dynamics folks, it would be great to get your opinion on how to get the second order gradient (assuming Phil could confirm he does need the second order gradient.)