<commit_before>
#include <stdexcept>
#include "RigidBodySystem.h"
#include "RigidBodyIK.h"  // required for resolving initial conditions
#include "urdfParsingUtil.h"

using namespace std;
using namespace Eigen;
using namespace Drake;


RigidBodySystem::StateVector<double> RigidBodySystem::dynamics(const double& t, const RigidBodySystem::StateVector<double>& x, const RigidBodySystem::InputVector<double>& u) const {
  using namespace std;
  using namespace Eigen;
  eigen_aligned_unordered_map<const RigidBody *, Matrix<double, 6, 1> > f_ext;

  // todo: make kinematics cache once and re-use it (but have to make one per type)
  auto nq = tree->num_positions;
  auto nv = tree->num_velocities;
  auto num_actuators = tree->actuators.size();
  auto q = x.topRows(nq);
  auto v = x.bottomRows(nv);
  auto kinsol = tree->doKinematics(q,v);

  // todo: preallocate the optimization problem and constraints, and simply update them then solve on each function eval.
  // happily, this clunkier version seems fast enough for now
  // the optimization framework should support this (though it has not been tested thoroughly yet)
  OptimizationProblem prog;
  auto const & vdot = prog.addContinuousVariables(nv,"vdot");

  auto H = tree->massMatrix(kinsol);
  Eigen::MatrixXd H_and_neg_JT = H;

  { // loop through rigid body force elements and accumulate f_ext

    // todo: distinguish between AppliedForce and ConstraintForce
    // todo: have AppliedForce output tau (instead of f_ext).  it's more direct, and just as easy to compute.

    int u_index = 0;
    const NullVector<double> force_state;  // todo:  will have to handle this case
    for (auto const & prop : props) {
      RigidBodyFrame* frame = prop->getFrame();
      RigidBody* body = frame->body.get();
      int num_inputs = 1;  // todo: generalize this
      RigidBodyPropellor::InputVector<double> u_i(u.middleRows(u_index,num_inputs));
      // todo: push the frame to body transform into the dynamicsBias method?
      Matrix<double,6,1> f_ext_i = transformSpatialForce(frame->transform_to_body,prop->output(t,force_state,u_i,kinsol));
      if (f_ext.find(body)==f_ext.end()) f_ext[body] = f_ext_i;
      else f_ext[body] = f_ext[body]+f_ext_i;
      u_index += num_inputs;
    }
  }

  VectorXd C = tree->dynamicsBiasTerm(kinsol,f_ext);
  if (num_actuators > 0) C -= tree->B*u.topRows(num_actuators);

  { // apply contact forces
    const bool use_multi_contact = false;
    VectorXd phi;
    Matrix3Xd normal, xA, xB;
    vector<int> bodyA_idx, bodyB_idx;
    if (use_multi_contact)
      tree->potentialCollisions(kinsol,phi,normal,xA,xB,bodyA_idx,bodyB_idx);
    else
      tree->collisionDetect(kinsol,phi,normal,xA,xB,bodyA_idx,bodyB_idx);

    for (int i=0; i<phi.rows(); i++) {
      if (phi(i)<0.0) { // then i have contact
        double mu = 1.0; // todo: make this a parameter

        // todo: move this entire block to a shared an updated "contactJacobian" method in RigidBodyTree
        auto JA = tree->forwardKinJacobian(kinsol, xA.col(i), bodyA_idx[i], 0, 0, false);
        auto JB = tree->forwardKinJacobian(kinsol, xB.col(i), bodyB_idx[i], 0, 0, false);
        Vector3d this_normal = normal.col(i);

        // compute the surface tangent basis
        Vector3d tangent1;
        if (1.0 - this_normal(2) < EPSILON) { // handle the unit-normal case (since it's unit length, just check z)
          tangent1 << 1.0, 0.0, 0.0;
        } else if (1 + this_normal(2) < EPSILON) {
          tangent1 << -1.0, 0.0, 0.0;  //same for the reflected case
        } else {// now the general case
          tangent1 << this_normal(1), -this_normal(0), 0.0;
          tangent1 /= sqrt(this_normal(1)*this_normal(1) + this_normal(0)*this_normal(0));
        }
        Vector3d tangent2 = tangent1.cross(this_normal);
        Matrix3d R;  // rotation into normal coordinates
        R << tangent1, tangent2, this_normal;
        auto J = R*(JA-JB);  // J = [ D1; D2; n ]
        auto relative_velocity = J*v;  // [ tangent1dot; tangent2dot; phidot ]

        if (false) {
          // spring law for normal force:  fA_normal = -k*phi - b*phidot
          // and damping for tangential force:  fA_tangent = -b*tangentdot (bounded by the friction cone)
          double k = 150, b = k / 10;  // todo: put these somewhere better... or make them parameters?
          Vector3d fA;
          fA(2) = -k * phi(i) - b * relative_velocity(2);
          fA.head(2) = -std::min(b, mu*fA(2)/(relative_velocity.head(2).norm()+EPSILON)) * relative_velocity.head(2);  // epsilon to avoid divide by zero

          // equal and opposite: fB = -fA.
          // tau = (R*JA)^T fA + (R*JB)^T fB = J^T fA
          C -= J.transpose()*fA;
        } else { // linear acceleration constraints (more expensive, but less tuning required for robot mass, etc)
          // note: this does not work for the multi-contact case (it overly constrains the motion of the link).  Perhaps if I made them inequality constraints...
          static_assert(!use_multi_contact, "The acceleration contact constraints do not play well with multi-contact");

          // phiddot = -2*alpha*phidot - alpha^2*phi   // critically damped response
          // tangential_velocity_dot = -2*alpha*tangential_velocity
          double alpha = 20;  // todo: put this somewhere better... or make them parameters?
          Vector3d desired_relative_acceleration = -2*alpha*relative_velocity;
          desired_relative_acceleration(2) += -alpha*alpha*phi(i);
          // relative_acceleration = J*vdot + R*(JAdotv - JBdotv) // uses the standard dnormal/dq = 0 assumption

          cout << "phi = " << phi << endl;
          cout << "desired acceleration = " << desired_relative_acceleration.transpose() << endl;
//          cout << "acceleration = " << (J*vdot + R*(JAdotv - JBdotv)).transpose() << endl;

          prog.addContinuousVariables(3,"contact normal force");
          auto JAdotv = tree->forwardJacDotTimesV(kinsol, xA.col(i).eval(), bodyA_idx[i], 0, 0);
          auto JBdotv = tree->forwardJacDotTimesV(kinsol, xB.col(i).eval(), bodyB_idx[i], 0, 0);

          prog.addLinearEqualityConstraint(J,desired_relative_acceleration - R*(JAdotv - JBdotv),{vdot});
          H_and_neg_JT.conservativeResize(NoChange,H_and_neg_JT.cols()+3);
          H_and_neg_JT.rightCols(3) = -J.transpose();
        }
      }
    }
  }

  if (tree->getNumPositionConstraints()) {
    int nc = tree->getNumPositionConstraints();
    const double alpha = 5.0;  // 1/time constant of position constraint satisfaction (see my latex rigid body notes)

    prog.addContinuousVariables(nc,"position constraint force");  // don't actually need to use the decision variable reference that would be returned...

    // then compute the constraint force
    auto phi = tree->positionConstraints(kinsol);
    auto J = tree->positionConstraintsJacobian(kinsol,false);
    auto Jdotv = tree->positionConstraintsJacDotTimesV(kinsol);

    // phiddot = -2 alpha phidot - alpha^2 phi  (0 + critically damped stabilization term)
    prog.addLinearEqualityConstraint(J,-(Jdotv + 2*alpha*J*v + alpha*alpha*phi),{vdot});
    H_and_neg_JT.conservativeResize(NoChange,H_and_neg_JT.cols()+J.rows());
    H_and_neg_JT.rightCols(J.rows()) = -J.transpose();
  }

  // add [H,-J^T]*[vdot;f] = -C
  prog.addLinearEqualityConstraint(H_and_neg_JT,-C);

  prog.solve();<commit_msg>Fix RigidBodySystem.<commit_after>
#include <stdexcept>
#include "RigidBodySystem.h"
#include "RigidBodyIK.h"  // required for resolving initial conditions
#include "urdfParsingUtil.h"

using namespace std;
using namespace Eigen;
using namespace Drake;


RigidBodySystem::StateVector<double> RigidBodySystem::dynamics(const double& t, const RigidBodySystem::StateVector<double>& x, const RigidBodySystem::InputVector<double>& u) const {
  using namespace std;
  using namespace Eigen;
  eigen_aligned_unordered_map<const RigidBody *, Matrix<double, 6, 1> > f_ext;

  // todo: make kinematics cache once and re-use it (but have to make one per type)
  auto nq = tree->num_positions;
  auto nv = tree->num_velocities;
  auto num_actuators = tree->actuators.size();
  auto q = x.topRows(nq);
  auto v = x.bottomRows(nv);
  auto kinsol = tree->doKinematics(q,v);

  // todo: preallocate the optimization problem and constraints, and simply update them then solve on each function eval.
  // happily, this clunkier version seems fast enough for now
  // the optimization framework should support this (though it has not been tested thoroughly yet)
  OptimizationProblem prog;
  auto const & vdot = prog.addContinuousVariables(nv,"vdot");

  auto H = tree->massMatrix(kinsol);
  Eigen::MatrixXd H_and_neg_JT = H;

  { // loop through rigid body force elements and accumulate f_ext

    // todo: distinguish between AppliedForce and ConstraintForce
    // todo: have AppliedForce output tau (instead of f_ext).  it's more direct, and just as easy to compute.

    int u_index = 0;
    const NullVector<double> force_state;  // todo:  will have to handle this case
    for (auto const & prop : props) {
      RigidBodyFrame* frame = prop->getFrame();
      RigidBody* body = frame->body.get();
      int num_inputs = 1;  // todo: generalize this
      RigidBodyPropellor::InputVector<double> u_i(u.middleRows(u_index,num_inputs));
      // todo: push the frame to body transform into the dynamicsBias method?
      Matrix<double,6,1> f_ext_i = transformSpatialForce(frame->transform_to_body,prop->output(t,force_state,u_i,kinsol));
      if (f_ext.find(body)==f_ext.end()) f_ext[body] = f_ext_i;
      else f_ext[body] = f_ext[body]+f_ext_i;
      u_index += num_inputs;
    }
  }

  VectorXd C = tree->dynamicsBiasTerm(kinsol,f_ext);
  if (num_actuators > 0) C -= tree->B*u.topRows(num_actuators);

  { // apply contact forces
    const bool use_multi_contact = false;
    VectorXd phi;
    Matrix3Xd normal, xA, xB;
    vector<int> bodyA_idx, bodyB_idx;
    if (use_multi_contact)
      tree->potentialCollisions(kinsol,phi,normal,xA,xB,bodyA_idx,bodyB_idx);
    else
      tree->collisionDetect(kinsol,phi,normal,xA,xB,bodyA_idx,bodyB_idx);

    for (int i=0; i<phi.rows(); i++) {
      if (phi(i)<0.0) { // then i have contact
        double mu = 1.0; // todo: make this a parameter

        // todo: move this entire block to a shared an updated "contactJacobian" method in RigidBodyTree
        auto JA = tree->transformPointsJacobian(kinsol, xA.col(i), bodyA_idx[i], 0, false);
        auto JB = tree->transformPointsJacobian(kinsol, xB.col(i), bodyB_idx[i], 0, false);
        Vector3d this_normal = normal.col(i);

        // compute the surface tangent basis
        Vector3d tangent1;
        if (1.0 - this_normal(2) < EPSILON) { // handle the unit-normal case (since it's unit length, just check z)
          tangent1 << 1.0, 0.0, 0.0;
        } else if (1 + this_normal(2) < EPSILON) {
          tangent1 << -1.0, 0.0, 0.0;  //same for the reflected case
        } else {// now the general case
          tangent1 << this_normal(1), -this_normal(0), 0.0;
          tangent1 /= sqrt(this_normal(1)*this_normal(1) + this_normal(0)*this_normal(0));
        }
        Vector3d tangent2 = tangent1.cross(this_normal);
        Matrix3d R;  // rotation into normal coordinates
        R << tangent1, tangent2, this_normal;
        auto J = R*(JA-JB);  // J = [ D1; D2; n ]
        auto relative_velocity = J*v;  // [ tangent1dot; tangent2dot; phidot ]

        if (false) {
          // spring law for normal force:  fA_normal = -k*phi - b*phidot
          // and damping for tangential force:  fA_tangent = -b*tangentdot (bounded by the friction cone)
          double k = 150, b = k / 10;  // todo: put these somewhere better... or make them parameters?
          Vector3d fA;
          fA(2) = -k * phi(i) - b * relative_velocity(2);
          fA.head(2) = -std::min(b, mu*fA(2)/(relative_velocity.head(2).norm()+EPSILON)) * relative_velocity.head(2);  // epsilon to avoid divide by zero

          // equal and opposite: fB = -fA.
          // tau = (R*JA)^T fA + (R*JB)^T fB = J^T fA
          C -= J.transpose()*fA;
        } else { // linear acceleration constraints (more expensive, but less tuning required for robot mass, etc)
          // note: this does not work for the multi-contact case (it overly constrains the motion of the link).  Perhaps if I made them inequality constraints...
          static_assert(!use_multi_contact, "The acceleration contact constraints do not play well with multi-contact");

          // phiddot = -2*alpha*phidot - alpha^2*phi   // critically damped response
          // tangential_velocity_dot = -2*alpha*tangential_velocity
          double alpha = 20;  // todo: put this somewhere better... or make them parameters?
          Vector3d desired_relative_acceleration = -2*alpha*relative_velocity;
          desired_relative_acceleration(2) += -alpha*alpha*phi(i);
          // relative_acceleration = J*vdot + R*(JAdotv - JBdotv) // uses the standard dnormal/dq = 0 assumption

          cout << "phi = " << phi << endl;
          cout << "desired acceleration = " << desired_relative_acceleration.transpose() << endl;
//          cout << "acceleration = " << (J*vdot + R*(JAdotv - JBdotv)).transpose() << endl;

          prog.addContinuousVariables(3,"contact normal force");
          auto JAdotv = tree->transformPointsJacobianDotTimesV(kinsol, xA.col(i).eval(), bodyA_idx[i], 0);
          auto JBdotv = tree->transformPointsJacobianDotTimesV(kinsol, xB.col(i).eval(), bodyB_idx[i], 0);

          prog.addLinearEqualityConstraint(J,desired_relative_acceleration - R*(JAdotv - JBdotv),{vdot});
          H_and_neg_JT.conservativeResize(NoChange,H_and_neg_JT.cols()+3);
          H_and_neg_JT.rightCols(3) = -J.transpose();
        }
      }
    }
  }

  if (tree->getNumPositionConstraints()) {
    int nc = tree->getNumPositionConstraints();
    const double alpha = 5.0;  // 1/time constant of position constraint satisfaction (see my latex rigid body notes)

    prog.addContinuousVariables(nc,"position constraint force");  // don't actually need to use the decision variable reference that would be returned...

    // then compute the constraint force
    auto phi = tree->positionConstraints(kinsol);
    auto J = tree->positionConstraintsJacobian(kinsol,false);
    auto Jdotv = tree->positionConstraintsJacDotTimesV(kinsol);

    // phiddot = -2 alpha phidot - alpha^2 phi  (0 + critically damped stabilization term)
    prog.addLinearEqualityConstraint(J,-(Jdotv + 2*alpha*J*v + alpha*alpha*phi),{vdot});
    H_and_neg_JT.conservativeResize(NoChange,H_and_neg_JT.cols()+J.rows());
    H_and_neg_JT.rightCols(J.rows()) = -J.transpose();
  }

  // add [H,-J^T]*[vdot;f] = -C
  prog.addLinearEqualityConstraint(H_and_neg_JT,-C);

  prog.solve();
{'fraction_non_alphanumeric': 0.12226808309087218, 'fraction_numerical': 0.01245009811218621, 'mean_word_length': 3.202445265851578, 'pattern_counts': {'":': 0, '<': 83, '<?xml version=': 0, '>': 69, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4319499', 'n_tokens_mistral': 5086, 'n_tokens_neox': 4796, 'n_words': 1276}