Closed diegoferigo closed 2 months ago
This is basically the updated version of the problem occurring in the previous OOP APIs in #84 cc @flferretti
The example is failing with the following error for me:
traversaro@IITBMP014LW012:~/jaxsimws$ pixi run python mwe103.py
jaxsim[21436] INFO Enabling JAX to use 64bit precision
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Traceback (most recent call last):
File "/home/traversaro/jaxsimws/mwe103.py", line 18, in <module>
urdf_string = rod.urdf.exporter.UrdfExporter.sdf_to_urdf_string(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/traversaro/jaxsimws/.pixi/envs/default/lib/python3.12/site-packages/rod/urdf/exporter.py", line 30, in sdf_to_urdf_string
model = sdf.models()[0]
~~~~~~~~~~~~^^^
IndexError: list index out of range
Installed packages:
Perhaps a specific rod version is required?
Also #83 actually needed https://github.com/ami-iit/rod/pull/27 for proper color support, we need to bump the dependencies before releasing jaxsim v0.2.0
.
Playing a bit around, I was able to reduce the MWE, to remove jaxsim at all:
import dataclasses
import jax
import jax.numpy as jnp
import jax_dataclasses
from jax_dataclasses import Static
import pathlib
import numpy.typing as npt
import dataclasses
from typing import Dict, Union
import jax.lax
import jax.numpy as jnp
import jax_dataclasses
import numpy as np
from jax_dataclasses import Static
@jax_dataclasses.pytree_dataclass
class MWE103GroundContact:
body: Static[npt.NDArray] = dataclasses.field(
default_factory=lambda: np.array([])
)
@staticmethod
def build_from() -> "MWE103GroundContact":
# Build the object
gc = MWE103GroundContact(body=np.zeros(8))
return gc
def local_total_mass(model: MWE103GroundContact) -> float:
return 1.0
dummy_contact1 = MWE103GroundContact.build_from()
dummy_contact2 = MWE103GroundContact.build_from()
local_total_mass_jit1 = jax.jit(local_total_mass)
local_total_mass_jit2 = jax.jit(local_total_mass)
_ = local_total_mass_jit1(model=dummy_contact1)
_ = local_total_mass_jit2(model=dummy_contact2)
Somehow the problem is related to the body
attribute of the GroundContact
class: https://github.com/ami-iit/jaxsim/blob/f3d9abb0dd8583c5391db7b47c998ad0a457adea/src/jaxsim/physics/model/ground_contact.py#L26 . Commenting it, the problem disappears.
Somehow the problem is related to the
body
attribute of theGroundContact
class:
This is aligned with what was found in https://github.com/ami-iit/jaxsim/issues/84#issuecomment-1943719169 .
I think the answer to this question clarifies everything: https://github.com/google/jax/issues/19547#issuecomment-1913711588 .
I think I described the root issue behind all the ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
in https://github.com/ami-iit/jaxsim/pull/105 .
Doing some additional tests, I found out that a similar problem is raised. This time, it seems to be related to the integrators:
Traceback (most recent call last):
File "/home/flferretti/git/comodo/src/comodo/jaxsimSimulator/test.py", line 55, in <module>
integrator_state2 = integrator2.init(x0=data2.state, t0=0, dt=1e-3)
File "/home/flferretti/jaxsim/src/jaxsim/integrators/common.py", line 166, in init
_ = integrator(x0, t0, dt, **kwargs)
File "/home/flferretti/jaxsim/src/jaxsim/integrators/common.py", line 283, in __call__
z = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
File "/home/flferretti/jaxsim/src/jaxsim/integrators/common.py", line 417, in _compute_next_state
K, _ = jax.lax.scan(
File "/home/flferretti/jaxsim/src/jaxsim/integrators/common.py", line 403, in scan_body
ki = jax.lax.cond(
File "/home/flferretti/jaxsim/src/jaxsim/integrators/common.py", line 374, in <lambda>
get_ẋ0 = lambda: self.params.get("dxdt0", f(x0, t0)[0])
File "/home/flferretti/jaxsim/src/jaxsim/integrators/common.py", line 365, in <lambda>
f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
File "/home/flferretti/jaxsim/src/jaxsim/api/ode.py", line 69, in f
return system_dynamics(
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[1]..
The error occurred while tracing the function <lambda> at /home/flferretti/jaxsim/src/jaxsim/integrators/common.py:374 for cond. This value became a tracer due to JAX operations on these lines:
operation a:bool[1] = eq b c
from line /home/flferretti/jaxsim/src/jaxsim/api/ode.py:69 (f)
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError
After some testing, I found out that the objects data.state.physics_model
are not compatible. Moreover, when data
is created for a model without joints, like in the example in https://github.com/ami-iit/jaxsim/issues/103#issuecomment-2058752595, the fact that the attributes joint_positions
and joint_velocities
are empty raises https://github.com/google/jax/blob/95c05521b457f2572b0704449c76982c85ee72b6/jax/_src/core.py#L646-L647.
Setting hash=False
for both attributes partially solves the problem, but still I cannot compare data1
with data2
to get a boolean
Furthermore, I believe that we should pay attention to the dynamics
attribute of Integrator
, which being just a Protocol, it cannot be compared with ==
Thanks for the investigation @flferretti. Since JaxSimModelData
contains many JAX arrays as attributes, I believe that the __hash__
and __eq__
methods need to be properly implemented. I'll have a look at it, I've already have some WIP work on the implementation of these methods for the jaxsim.parser
package, I'll extend it also to JaxSimModelData
.
Also with the new functional APIs, the following problem still affects our code:
A MWE to reproduce the problem is the following:
I've already tried to play a bit with
JaxSimModel.__eq__
andJaxSimModel.__hash__
with no luck. Marking this issue as help wanted in case this problem becomes annoying to someone.Interestingly, running it in a context in which
jax.jit
is disabled, works fine:I suspect that there is something strange in the jax cache that we are either overlooking or not handling properly.