Refactor contact forces sum in api.ode.system_velocity_dynamics

[flferretti]

This PR refactors contact forces sum in system_velocity_dynamics potentially improving readability and performance

---

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

[diegoferigo]

What was the original desiderata of this PR? Runtime performance, memory footprint, readability, etc?

It seems to me that now readability is much worse than before, and on a quick test I've done, the new logic seems wrong. Did you double check it?

[flferretti]

The idea was both to make the code simpler, as the current version is quite intricate, and reduce the memory footprint of that part of code. If you prefer, I can split the binary mask and the jnp.where(...) as:

Suggested change

	W_f_Li_terrain = jnp.where(
	(
	parent_link_index_of_collidable_points[:, None]
	== jnp.arange(model.number_of_links())
	).any(axis=-1, keepdims=True),
	W_f_Ci,
	jnp.zeros_like(W_f_Ci),
	).sum(axis=0)
	mask = (
	parent_link_index_of_collidable_points[:, None]
	== jnp.arange(model.number_of_links())
	).any(axis=-1, keepdims=True)
	W_f_Li_terrain = jnp.where(
	mask,
	W_f_Ci,
	jnp.zeros_like(W_f_Ci),
	).sum(axis=0)

[flferretti]

What was the original desiderata of this PR? Runtime performance, memory footprint, readability, etc?

It seems to me that now readability is much worse than before, and on a quick test I've done, the new logic seems wrong. Did you double check it?

I answered you with a comment in the code.

For what regards instead the logic, in which example did you get the error? I did the test simulating spheres in vmap and I had no problem

[diegoferigo]

For what regards instead the logic, in which example did you get the error? I did the test simulating spheres in vmap and I had no problem

Check the following snippet, I get different results with the new logic. Did I do any mistake?

Code

import jax
import jax.numpy as jnp
import jaxsim.api as js
import resolve_robotics_uri_py
import rod

# Find the urdf file.
urdf_path = resolve_robotics_uri_py.resolve_robotics_uri(
    uri="model://ergoCubSN001/model.urdf"
)

# Build the ROD model.
rod_sdf = rod.Sdf.load(sdf=urdf_path)

# Build the model.
model = js.model.JaxSimModel.build_from_model_description(
    model_description=rod_sdf.model,
)

# Get the parent body of the collidable points.
parent_link_index_of_collidable_points = jnp.array(
    model.kin_dyn_parameters.contact_parameters.body
)

# Create the 6D forces of collidable points.
nc = len(parent_link_index_of_collidable_points)
W_f_Ci = jnp.ones(shape=(nc, 6))

# =========
# Old logic
# =========

f_old = jax.vmap(
    lambda nc: (
        jnp.vstack(jnp.equal(parent_link_index_of_collidable_points, nc).astype(int))
        * W_f_Ci
    ).sum(axis=0)
)(jnp.arange(model.number_of_links()))

# =========
# New logic
# =========

f_new = jnp.where(
    (
        parent_link_index_of_collidable_points[:, None]
        == jnp.arange(model.number_of_links())
    ).any(axis=-1, keepdims=True),
    W_f_Ci,
    jnp.zeros_like(W_f_Ci),
).sum(axis=0)

[flferretti]

I found another way to further simplify the code, if you prefer I could merge the two lines. Ready for review @diegoferigo, thanks!

[diegoferigo]

Much better. I also tested on the script of my previous comment and it seems working fine now.