Skip to content

3. How the code is organized

Jump to bottom
Michelle edited this page Oct 30, 2021 · 8 revisions

Entrypoints

The main entrypoint for a physics timestep is dart::simulation::World::step(). This is responsible for computing a single timestep of a physical simulation on a World object.

A World is composed of a bunch of dart::dynamics::Skeleton's, each of which has lots of dart::dynamics::Joint's connecting dart::dynamics::BodyNode's. Every Skeleton's root Joint is assumed to be fixed to the origin of the world. So if you want your Skeleton to move around, you need to attach a dart::dynamics::FreeJoint as the root joint, which attaches the root BodyNode of your Skeleton to the World.

In Nimble, in order to save all the information we need in order to compute gradients backwards through a physics timestep, we wrap the call to dart::simulation::World::step() inside of dart::neural::forwardPass(world). You can find the implementation of this method in dart/neural/NeuralUtils.cpp. This takes a forward step in the world, but also constructs and returns a dart::neural::BackpropSnapshot object, which saves all the information within itself to compute gradients.

The key entry points within dart::neural::BackpropSnapshot are all the getXXJacobian() methods:

  • dart::neural::BackpropSnapshot::getVelVelJacobian(): This computes the Jacobian relating the the velocity at the beginning of this timestep v_t to the velocity at the end of the timestep v_t+1
  • dart::neural::BackpropSnapshot::getVelPosJacobian(): This computes the Jacobian relating the the velocity at the beginning of this timestep v_t to the position at the end of the timestep p_t+1
  • dart::neural::BackpropSnapshot::getPosVelJacobian(): This computes the Jacobian relating the the position at the beginning of this timestep p_t to the velocity at the end of the timestep v_t+1
  • dart::neural::BackpropSnapshot::getPosPosJacobian(): This computes the Jacobian relating the the position at the beginning of this timestep p_t to the position at the end of the timestep p_t+1
  • dart::neural::BackpropSnapshot::getForceVelJacobian(): This computes the Jacobian relating the the force at the beginning of this timestep f_t to the velocity at the end of the timestep v_t+1

These methods form the backbone of the DiffDART implementation. Everything else to do with differentiable physics is built on top of the Jacobians returned by these methods.

How forward physics computation works

When you call World::step(), a bunch of stuff happens in order to compute a physics timestep. Here's what it is, and where it happens:

Step 1.

We call Skeleton::computeForwardDynamics() and then Skeleton::integrateVelocities() on each skeleton. The computeForwardDynamics() call takes the current positions, velocities, and forces on that Skeleton and compute acceleration (for now, ignoring collisions). Then Skeleton::integrateVelocities() basically does joint_velocity[i] += joint_acceleration[i] * time_step. It's a bit more complex than just simple addition because FreeJoint and BallJoint have special treatment. Basically, we've just run a normal physics timestep if there were no collisions. In matrix notation, this step is:

v_t+1 = v_t + dt * M^{-1} * (C(p_t, v_t) + tau_t)

Where C(p_t, v_t) is the Coriolis and gravity force in joint space at a given orientation p_t and velocity v_t.

Step 2.

The next important line in World::step() is mConstraintSolver->solve(this). This deals with contact. Let's break this down into pieces:

Step 2.1.

mConstraintSolver->solve(this) first runs collision detection, which ends up calling dart::collision::dart::DARTCollide::collide() on every pair of objects to check if any are intersecting. You can find that in dart/collision/dart/DARTCollide.cpp. Any collisions it finds get annotated with a ton of meta-data, including a ContactType enum. ContactType is defined in dart/collision/Contact.hpp, and covers all sorts of cases like ContactType::VERTEX_FACE and ContactType::EDGE_EDGE, ContactType::SPHERE_FACE, etc. This information is useful later for computing gradients of how contact points and normals change. If there's friction present, each contact will become 3 possible constraint forces, one for the normal force, and two mutually orthogonal vectors for friction. Without friction, it's just the single normal force per contact.

Step 2.2.

Once mConstraintSolver->solve(this) has found the collisions and all the contact force positions and normals (and friction bounds), it sets up a Linear Complimentarity Problem (LCP) and solves it in order to find contact impulses to prevent all the contacts from inter-penetrating. This happens in dart/constraint/BoxedLcpConstraintSolver.cpp, in the method BoxedLcpConstraintSolver::solveConstrainedGroup(). That method first sets up a matrix A and a vector b, and then solves f = LCP(A,b) to find the constraint impulses f, which has one entry per contact constraint. The vector b holds the relative velocity of each contact point (positive means it's separating, negative means it's moving together). The vector A holds something like a mass matrix, but in contact space. If you apply a vector of constraint impulses f, you'll get a change in contact point relative velocity delta_v = A*f.

The LCP solver also sets up "friction bounds" on the LCP, to ensure that no frictional force is greater than the magnitude of the corresponding contact-normal force times the friction coefficient. The vector fIndex holds pointers to which index a given force is "dependent on" for friction. If fIndex[5] = 3, then we have the bound that f[3] * lo[5] <= f[5] <= f[3] * hi[5]. If findex[i] = -1, then we say that this is a normal force, and so isn't dependent on the magnitude of any other forces: lo[i] <= f[i] <= hi[i].

Because of the friction bounds, the LCP isn't guaranteed to be solvable. It usually is solvable, but sometimes the solvers fail, so be aware that it can happen.

Note:

As an optimization, DART sometimes breaks up a big LCP into a number of smaller LCPs for a given timestep. Since LCP solving is at least O(n^3) runtime, several small LCPs is faster than one big one. DART sets up multiple "collision groups", one for each set of independently colliding objects (example: blocks A and B are colliding, so in a group together, but maybe block C is far away from both, so doesn't need to be included for collision solving LCP), to reduce the size of each LCP solve. So each "collision group" basically gets treated as its own little world for the forward pass. These collision groups are called ConstrainedGroups and each generate a data log for gradients called ConstrainedGroupGradientMatrices. Our BackpropSnapshot actually contains a list of ConstrainedGroupGradientMatrices, and constructs matrices for the whole world by aggregating results from all the sub-clusters of collisions.

Step 2.3.

After we've found a vector of forces f, at the very end of BoxedLcpConstraintSolver::solveConstrainedGroup() we call constraint->applyImpulse() and constraint->excite() for each constraint, which gets the impulses locked in and ready to be computed.

Step 3.

Popping the stack all the way back to World::step(), we now apply our impulses. We do this in the block:

if (skel->isImpulseApplied())
{
  skel->computeImpulseForwardDynamics();
  skel->setImpulseApplied(false);
}

In matrix math, this is equivalent to:

v_t+1 += J*f

Where J is some Jacobian relating constraint forces to joint velocity accelerations. In the code, the J matrix is broken up into columns based on whether the indices are clamping or not. For the clamping columns of J (which for historical reasons we often refer to as A_c) you can get that matrix by calling BackpropSnapshot::getClampingConstraintMatrix() on a snapshot.

Step 4.0:

Last but not least, now that we've finally computed v_t+1 (first by taking an unconstrained velocity integration, then bounding it with the LCP impulses), it's time to integrate position and compute p_t+1.

There are fundamentally two ways to do this, and it's controlled by the mParallelVelocityAndPositionUpdates flag on World. The first (classic) way to update position, when mParallelVelocityAndPositionUpdates = false, is with this formula:

p_t+1 = p_t + dt*v_t+1

Using v_t+1 in this integration introduces headaches for backpropagation through time, because now p_t+1 depends on v_t+1, so the standard practice of concatenating them together into a single state vector won't work. This breaks a lot of classic control theory stuff, so we support another position integration scheme. When mParallelVelocityAndPositionUpdates = true we break the dependence between p_t+1 and v_t+1 by changing the integration scheme to the following:

p_t+1 = p_t + dt*v_t

This makes it possible to concatenate p_t+1 and v_t+1 in a single vector without breaking Jacobians, but it introduces some small instability in forward simulation.

Backprop:

And with that, we've computed an entire forward pass. All that's left is the main trick of DiffDART, which is relating the vectors p_t+1 and v_t+1 to the vectors p_t, v_t, and tau_t. Those Jacobians, computable on BackpropSnapshot, fill out the following table:

p_t+1 v_t+1
p_t getPosPosJacobian() getPosVelJacobian()
v_t getVelPosJacobian() getVelVelJacobian()
tau_t Always 0 getForceVelJacobian()

In order to do that, we'll need to go over what the formula for the forward relationship in full. Here it is:

v_t+1 = v_t + dt * M^{-1} * (C(p_t, v_t) + tau_t) + J * LCP(A,b)

One of the core contributions of the paper is to note that with small enough perturbations around the previous solution, the LCP can be replaced with a linear equation. LCP(A,b) is the equivalent of Q^{-1} b_c for some square matrix Q (which we'll describe in a second) and the subset of b that corresponds to the clamping indices of b, which we'll call b_c. That gives us all the non-zero contact forces f_c (the clamping forces). All the other forces can be ignored.

In a confusing abuse of notation that exists in the code for historical reasons (it took a few drafts of the code to clarify my thinking on the math), we call the subset of J's columns that correspond to clamping indices A_c. This A_c doesn't have much have to do with the A we've been talking about with the LCP(A,b). I'm sorry about this, and we should rename this variable to something like J_c, but it's everywhere, so it's obnoxious to rewrite.

With A_c, which you can get with BackpropSnapshot::getClampingConstraintMatrix(world), we can be very explicit about our LCP substitution.

Q = A_c.transpose() * M^{-1} * A_c
b_c = -A_c.transpose() * (v_t + dt * M^{-1} * (C(p_t, v_t) + tau_t))
f_c = Q^{-1}b_c

And finally, rewriting our whole forward step in linear operations to get rid of the LCP:

v_t+1 = v_t + dt * M^{-1} * (C(p_t, v_t) + tau_t) + A_c.transpose() * f_c

Changes in v_t and tau_t change only b_c (and themselves, obviously). So we need to know how b_c changes with respect to v_t and tau_t. Since b_c is just the relative velocity at the contact points after one unincumbered forward step, we have:

b_c = - A_c.transpose() * (v_t + dt * M^{-1} * (C(p_t, v_t) + tau_t))

Then we can work out the Jacobian of b_c with respect to v_t and tau_t. That quantity can be gotten with BackpropSnapshot::getJacobianOfLCPOffsetClampingSubset(world, WithRespectTo::VELOCITY) and BackpropSnapshot::getJacobianOfLCPOffsetClampingSubset(world, WithRespectTo::FORCE). If we're getting the Jacobian of b_c with respect to C(p_t, v_t) we have to care about how the Coriolis forces change with velocity. That Jacobian can be gotten with BackpropSnapshot::getJacobianOfC(world, WithRespectTo::VELOCITY).

Once we have dB = BackpropSnapshot::getJacobianOfLCPOffsetClampingSubset(...) we can compute the higher level Jacobians getVelVelJacobian() and getForceVelJacobian() with a minimum of complexity.

Things get more complex when we're taking Jacobians with respect to p_t. Position touches everything. The mass matrix M^{-1} changes with position. So does the clamping columns of the contact Jacobian A_c. That means so does Q. Obviously also so does b_c. And so does the Coriolis force C(p_t, v_t). Those map as follows:

Variable Jacobian method
Q getJacobianOfLCPConstraintMatrixClampingSubset(world, b, wrt)
A_c getJacobianOfClampingConstraints(world, rhs)
A_c.transpose() getJacobianOfClampingConstraintsTranspose(world, rhs)
b_c getJacobianOfLCPOffsetClampingSubset(world, wrt)
M^{-1} getJacobianOfMinv(world, rhs, wrt)
C(p_t, v_t) getJacobianOfC(world, wrt)

The astute reader will be worried now that some of these methods taking gradients of matrices wrt vectors return Tensors, and that seems dimensionally messy. Fear not! All our methods return the tensor right-multiplied by a constant vector rhs, so that you still get an easy-to-work-with matrix. For example, getJacobianOfMinv(world, rhs, wrt) returns the result of differentiating the expression M^{-1}*rhs with respect to wrt, where rhs is a constant vector.

With those Jacobians in place, it's possible to compute getPosVelJacobian().

The last bits we need to worry about are getPosPosJacobian() and getVelPosJacobian(). Remember that our time integration scheme looks like this:

p_t+1 = p_t + dt*v_t

So it's easy to see that getPosPosJacobian() = I and getVelPosJacobian() = dt * I. There's some complexity when we have bounces, where we want our Jacobian to approximate a continuous time system, so that's what you see in the code. Those details are beyond the current scope of this document.

And with that, you should have all the info you need to go spelunking around in the code! God Speed, and you should continue on to Finding and Fixing Bugs.

Clone this wiki locally