Add new ViscoElasticContacts

[diegoferigo]

This PR implements a revised and ―possibly― improved contact model originally proposed in the following manuscript:

Exponential Integration for Efficient and Accurate Multi-Body Simulation with Stiff Viscoelastic Contacts
Hammoud, Olivieri, Righetti, Carpentier, Del Prete
https://arxiv.org/abs/2101.06846

Although the theory is not yet public, my implementation proposes the following improvements:

1. The original formulation only considered linear contact models. If $\delta$ and $\dot{\delta}$ are the penetration depth and the penetration rate along the terrain's normal direction, the new formulation enables using any non-linear contact model belonging to the family of Hunt/Crossley models:

$$f^{\perp} = (K \delta^p) \, \delta + (D \delta^q) \, \dot{\delta}$$

2. The new formulation allows to use non-linear contact models that intrinsically support uneven terrain, under the assumption that we can gather the terrain normal and the terrain height at any $x-y$ coordinate.
3. The new formulation does not integrate the contact dynamics in terms of displacement w.r.t. the anchor point $\mathbf{x} = [\Delta \mathbf{p}; \, \Delta \dot{\mathbf{p}}],$ but it integrates the contact dynamics ignoring the anchor point entirely, therefore $\mathbf{x} = [\mathbf{p}; \, \dot{\mathbf{p}}].$ This enables to do not keep track of the location where the contact was made.
4. The new formulation changes completely how the frictional component of the force (i.e., the tangential component) is computed. It moves the logic to compute the sticking tangential force inside the non-linear contact model, exploiting the fact that we can adopt AD to linearize the contact model with arbitrarily complex logic. The tangential displacement to obtain the $\Delta$ variables are computed by considering an additional state $\mathbf{m}$ of the terrain deformation (that, in a sense, replaces the need of keeping track of the contact creation). This workaround was inspired by the soft-contact model already implemented in the soft contacts of JaxSim. For this reason, the state of the contact dynamics becomes $\mathbf{x} = [\mathbf{p}; \, \dot{\mathbf{p}}; \, \mathbf{m}].$
5. The introduction of $\mathbf{m}$ in the state of the contact dynamics allows us to integrate this variable in continuous time, obtaining a stable and precise computation of the contact forces.
6. The new formulation, thanks to the presence of $\mathbf{m}$ in the state of the contact dynamics and the possibility to have arbitrarily complex non-linear contact models, completely changes the original implementation of sticking/slipping transition. I adapted the sticking/slipping logic used by the default soft contact model of JaxSim to this new contact model, and I obtained sound theoretical results (that also works well in practice, paying the price of a larger $A$ matrix used in $\exp^{At}$).

viscoelastic_contacts_ergocub.mp4

Open points left to future work:

• The main bottleneck of the whole contact model is the computation of the matrix exponential. In the paper, the authors describe a custom implementation that might be considerably faster than jax.scipy.linalg.expm. However, beyond implementing the logic, which is doable, we should also figure out how to expose AD. It's not clear to me what could be the performance degradation of relying to pure AD w.r.t. the usage of the custom adjoints as done internally in JAX.
• Exposing to the user the logic to enable/disable collidable points is not included in this PR. For the time being, altering directly the low-level boolean array is necessary for thinnin out the considered points.

---

📚 Documentation preview 📚: https://jaxsim--248.org.readthedocs.build//248/

[xela-95]

Hi @diegoferigo, another early comment since I'm testing this contact model: in jaxsim.api.ode.system_dynamics @

jaxsim/src/jaxsim/api/ode.py

Lines 378 to 387 in 986854e

	case SoftContacts():
	ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
	case RigidContacts() | RelaxedRigidContacts():
	pass
	case _:
	raise ValueError("Unable to determine contact state class prefix.")
	# Extract the velocities.

we should add a case for this class of contact model.

From collidable_point_dynamics I can see this model outputs some auxiliary data, that should be extracted:

jaxsim/src/jaxsim/api/contact.py

Lines 235 to 236 in 986854e

	aux_data = dict(W_f_avg2_C=W_f̿_Ci, m_tf=m_tf)

[diegoferigo]

[...]
we should add a case for this class of contact model.

From collidable_point_dynamics I can see this model outputs some auxiliary data, that should be extracted: [...]

Thanks @xela-95 for the feedback. The dictionary returned by the visco-elastic contact model is somewhat different then the one returned by the soft contact model. The latter returns $\dot{\mathbf{m}}$ (the derivative of the material deformation) that needs to be integrated together with the generalized state $(\mathbf{q}, \, \boldsymbol{\nu})$. The visco-elastic contacts, instead, directly compute $\mathbf{m}(t+\Delta t)$, therefore we need to override $\mathbf{m}$ of the extended state similarly than the velocity reset done in by the rigid contacts.

Regardless, I'm still unsure how to integrate the new visco-elastic contacts with the default jaxsim.model.step function. What I just wrote is only necessary when the default integrators are used. The best way to exploit this new contact model, however, is to use its own integrator jaxsim.rbda.contacts.visco_elastic.step, that integrates both the generalized position and velocity using the contact forces computed explicitly with the continuous exponential integrator.

[flferretti]

Thanks Diego, this is great! In general, would you agree in adding an underscore before the "private" methods so it is easier to understand what the user should not call externally?

[traversaro]

Regardless, I'm still unsure how to integrate the new visco-elastic contacts with the default jaxsim.model.step function. What I just wrote is only necessary when the default integrators are used. The best way to exploit this new contact model, however, is to use its own integrator jaxsim.rbda.contacts.visco_elastic.step, that integrates both the generalized position and velocity using the contact forces computed explicitly with the continuous exponential integrator.

Perhaps we can simply print a nice error if jaxsim.model.step is called with visco-elastic contacts ?

[diegoferigo]

Regardless, I'm still unsure how to integrate the new visco-elastic contacts with the default jaxsim.model.step function. What I just wrote is only necessary when the default integrators are used. The best way to exploit this new contact model, however, is to use its own integrator jaxsim.rbda.contacts.visco_elastic.step, that integrates both the generalized position and velocity using the contact forces computed explicitly with the continuous exponential integrator.

Perhaps we can simply print a nice error if jaxsim.model.step is called with visco-elastic contacts ?

After my comment, I already implemented a jit-compatible exception in this PR that is raised when the new contact model is used with the regular integrators.

jaxsim/src/jaxsim/api/contact.py

Lines 212 to 220 in 2bd727b

	# It is not yet clear how to pass the time step to this stage.
	# A possibility is to restrict the integrator to only forward Euler
	# and store the Δt inside the model.
	module = jaxsim.rbda.contacts.visco_elastic.step.__module__
	name = jaxsim.rbda.contacts.visco_elastic.step.__name__
	msg = "You need to use the custom '{}.{}' function with this contact model."
	jaxsim.exceptions.raise_runtime_error_if(
	condition=True, msg=msg.format(module, name)
	)

[diegoferigo]

@andreadelprete you might be interested in having a look at this PR. After we discussed about this methodology f2f, I was interested in drafting an implementation. In open-loop, it seems working pretty well. We are investigating if we can run our closed-loop control stack that currently works using both RigidContacts and RelaxedRigidContacts (as shown in #251).

Note that, as I wrote in the PR description, I updated the formulation to properly account for static friction and non-linear contact models. Performance are not yet the best due to the usage of the plain jax.scipy.linalg.expm, but it's a good starting point for playing with this integration method. Note also that we use our own estimation of the spring and damper parameters, that seemed to be quite effective also in this case.

Edit: last comment, the video reported above was recorded by running a simulation on GPU with 32-bits floats :)

[andreadelprete]

Hi @diegoferigo , thanks for pointing this out to me. Actually @DanielePucci has been bragging about this work with me for the last two weeks, so I was curious to see it! :)

3. The new formulation does not integrate the contact dynamics in terms of displacement w.r.t. the anchor point $\mathbf{x} = [\Delta \mathbf{p}; \, \Delta \dot{\mathbf{p}}],$ but it integrates the contact dynamics ignoring the anchor point entirely, therefore $\mathbf{x} = [\mathbf{p}; \, \dot{\mathbf{p}}].$ This enables to do not keep track of the location where the contact was made.

This is interesting, but I really don't get how you could achieve this. I guess I could understand by looking at the code, but that would take me a long time. Maybe you can explain me here instead?

4. The tangential displacement to obtain the $\Delta$ variables are computed by considering an additional state $\mathbf{m}$ of the terrain deformation (that, in a sense, replaces the need of keeping track of the contact creation). This workaround was inspired by the soft-contact model already implemented in the soft contacts of JaxSim. For this reason, the state of the contact dynamics becomes $\mathbf{x} = [\mathbf{p}; \, \dot{\mathbf{p}}; \, \mathbf{m}].$

Nice that you changed the tangential force computation. I was not completely satisfied by our results when slippage occured, so I expect you probably improved this.

I adapted the sticking/slipping logic used by the default soft contact model of JaxSim to this new contact model, and I obtained sound theoretical results (that also works well in practice, paying the price of a larger $A$ matrix used in $\exp^{At}$).

How much larger?

If one day you write the theory down, I'd love to have a look.

Andrea

[diegoferigo]

3. The new formulation does not integrate the contact dynamics in terms of displacement w.r.t. the anchor point $\mathbf{x} = [\Delta \mathbf{p}; , \Delta \dot{\mathbf{p}}],$ but it integrates the contact dynamics ignoring the anchor point entirely, therefore $\mathbf{x} = [\mathbf{p}; , \dot{\mathbf{p}}].$ This enables to do not keep track of the location where the contact was made.

This is interesting, but I really don't get how you could achieve this. I guess I could understand by looking at the code, but that would take me a long time. Maybe you can explain me here instead?

This is a workaround that we were already using in our simple SoftContacts model¹. It is also the methodology that enables the new formulation of sticking/slipping transitions.

The main intuition is that, instead of detecting/saving/updating $\mathbf{p}_0$, we introduce a new variable $\mathbf{m}$ that models the material deformation. Contrarily to $\mathbf{p}_0$, that is static and is updated only when a new contact is made, we consider $\mathbf{m}$ as a variable with its own dynamics $\dot{\mathbf{m}}$. Our non-linear contact model takes care to select the right dynamics of the material deformation depending on the contact status (off/sticking/slipping). You can grasp the intuition by considering that you can kind-of compute the displacement between the current position of the contact point w.r.t. $\mathbf{p}_0$ by introducing an extra variable in your state that integrates step-after-step the contact velocity. This is usually an approximation with discrete integrators, however with the exponential integration it becomes much more precise.

The original idea was taken from this paper², that I extended to non-flat surfaces in my PhD thesis³. The resulting system that allows for sticking/slipping can be seen as a spring/damper/clutch (Fig. 12 of the Azad paper):

How much larger?

We include this extra $\mathbf{m}$ variable in the system that is integrated with the matrix exponential. Therefore, the matrix $A$ of Eq. 11 of your paper is 50% larger, since the state vector in our case becomes $\mathbf{x} = [\mathbf{p}, \, \dot{\mathbf{p}}, \, \mathbf{m}]$.

Footnotes

1. This contact model operates independently on all contact points. Is does not consider them attached to a multibody system. ↩

2. Azad and Featherstone, Modeling the contact between a rolling sphere and a compliant ground plane, 2010, url. ↩

3. Ferigo, Chapter 7, Simulation Architectures for Reinforcement Learning applied to Robotics, 2022, url. ↩

[andreadelprete]

You can grasp the intuition by considering that you can kind-of compute the displacement between the current position of the contact point w.r.t. p0\mathbf{p}_0 by introducing an extra variable in your state that integrates step-after-step the contact velocity.

Ok, got it. So you get rid of p0 but you introduce m, so there is no gain in terms of quantity of information that you have to remember. The gain is regarding the sticking-slipping dynamics.

We include this extra m\mathbf{m} variable in the system that is integrated with the matrix exponential. Therefore, the matrix AA of Eq. 11 of your paper is 50% larger, since the state vector in our case becomes x=[p,p˙,m]\mathbf{x} = [\mathbf{p}, , \dot{\mathbf{p}}, , \mathbf{m}].

Ok. Unfortunately 50% larger is a significant increase, but if that's the price to pay to get a nice stick-slip behavior it's definitely worth it!

Thanks for the quick and clear reply!

Andrea

[diegoferigo]

Ok, got it. So you get rid of p0 but you introduce m, so there is no gain in terms of quantity of information that you have to remember. The gain is regarding the sticking-slipping dynamics.

Ok. Unfortunately 50% larger is a significant increase, but if that's the price to pay to get a nice stick-slip behavior it's definitely worth it!

Exactly, there is no real workaround to obtain sticking/slipping without somehow tracking the point displacement. You can find different ways to accumulate/compute this displacement, but if this information (corresponding to the tangential terrain direction) is used together with a second Hunt/Crossley-like model to compute the friction force, there's no really any solution I can think of that allows us to remove its need.

I agree that 50% more rows and columns is a significant increase, and this is the real trade-off of my proposed changes. In a sense, the formulation of this PR provides a self-contained solution without tricks that happen outside the contact model itself to obtain a behavior that intrinsically supports stick/slip transitions operating with a real friction cone (no pyramid or other approximations).

In a JAX-based context, this approach seems to be beneficial under many terms. I see this methodology as something that is more complex than our simple SoftContacts model (plain Hunt/Crossley model) since these visco-elastic contacts account for the actual robot dynamics projected to the contact space. There's more computation involved (Delassus matrix and exponential integration), but that's ok. The main benefit I see is that this exponential integration represents, similarly to SoftContacts an alternative contact model that does not rely on an iterative solver. This is really relevant for our context, as we are noticing that other contact models using either QP solvers (RigidContacts) or L-BFGS (RelaxedRigidContacts) currently run pretty slow and do not scale well with neither the number of contact points nor the number of DoFs.

Anyway, thanks for pointing me out to this methodology. I had fun studying the theory and implementing it, linearization and exponential integration brought me back to my time at university :)