ValueError is raised when using two models

[flferretti]

When calling the same method on two different high_level.Model instance, a ValueError is raised:

Minimal reproducible example:

from jaxsim.high_level.common import VelRepr
from jaxsim.high_level.model import Model

urdf_robot_string = ... # some path

model_1 = Model.build_from_model_description(
    model_description=urdf_robot_string,
    model_name="stickBot",
    vel_repr=VelRepr.Mixed,
    is_urdf=True,
).mutable(True, validate=False)

model_1.reset_base_position(position=jnp.array([0.1, -0.8, 0.5]))

model_2 = Model.build_from_model_description(
    model_description=urdf_robot_string,
    model_name="stickBot",
    vel_repr=VelRepr.Mixed,
    is_urdf=True,
).mutable(True, validate=False)

model_2.reset_base_position(jnp.array([1.0, 0.0, 0.0]))

Throws:

Traceback (most recent call last):
  File "/home/icub/git/test.py", line 65, in <module>
    model_2.reset_base_position(jnp.array([1.0, 0.0, 0.0]))
  File "/home/icub/git/jaxsim/src/jaxsim/utils/oop.py", line 189, in wrapper
    out, obj = fn_db(*args, **kwargs)
  File "/home/icub/git/jaxsim/src/jaxsim/utils/oop.py", line 400, in fn_tf_jit
    return jax.jit(function_to_jit, static_argnames=static_argnames)(
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I initially thought it was related to this jax-ml/jax#4717, but apparently it has been solved.

Nevertheless, another possibility is that the error is raised here, in jax_dataclasses as:

jax.tree_util.tree_structure(model_1) == jax.tree_util.tree_structure(model_2)

Throws:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

For some reason, it looks like the two models are not treated as two separate entities, and at some point the comparison between the models is performed inside jax_dataclasses

[flferretti]

Additional info:

• The output of jax.tree_util.tree_structure(...) is identical for the two objects
• PyTreeDef.__eq__ doesn't seem to be the problem
• Comparing children tree structures:

PyTreeDef(CustomNode(PhysicsModelState[()], [*, *, *, *, *, *])) == PyTreeDef(CustomNode(PhysicsModelState[()], [*, *, *, *, *, *]))
>>> True

PyTreeDef(CustomNode(GroundContact[(array([20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 33, 33, 33,
       33, 33, 33, 33, 33, 34, 34, 34, 34, 34, 34, 34, 34]),)], [*])) == PyTreeDef(CustomNode(GroundContact[(array([20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 33, 33, 33,
       33, 33, 33, 33, 33, 34, 34, 34, 34, 34, 34, 34, 34]),)], [*]))
>>> ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I verified that each children had the same PyTreeDef and I found out that the only issue is raised with the GroundContact branch