jaxsim 0.4.3.dev88__py3-none-any.whl → 0.4.3.dev90__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev88'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev88')
+__version__ = version = '0.4.3.dev90'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev90')

jaxsim/rbda/crba.py

@@ -94,48 +94,48 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
         j = i
 
-        CarryInnerFn = tuple[jtp.Int, jtp.Matrix, jtp.Matrix]
-        carry_inner_fn = (j, Fi, M)
+        FakeWhileCarry = tuple[jtp.Int, jtp.Vector, jtp.Matrix]
+        fake_while_carry = (j, Fi, M)
 
-        def while_loop_body(carry: CarryInnerFn) -> CarryInnerFn:
-            j, Fi, M = carry
+        # This internal for loop implements the while loop of the CRBA algorithm
+        # to compute off-diagonal blocks of the mass matrix M.
+        # In pseudocode it is implemented as a while loop. However, in order to enable
+        # applying reverse-mode AD, we implement it as a nested for loop with a fixed
+        # number of iterations and a branching model to skip for loop iterations.
+        def fake_while_loop(
+            carry: FakeWhileCarry, i: jtp.Int
+        ) -> tuple[FakeWhileCarry, None]:
 
-            Fi = i_X_λi[j].T @ Fi
-            j = λ[j]
-            jj = j - 1
+            def compute(carry: FakeWhileCarry) -> FakeWhileCarry:
 
-            M_ij = Fi.T @ S[j]
+                j, Fi, M = carry
 
-            M = M.at[ii + 6, jj + 6].set(M_ij.squeeze())
-            M = M.at[jj + 6, ii + 6].set(M_ij.squeeze())
+                Fi = i_X_λi[j].T @ Fi
+                j = λ[j]
 
-            return j, Fi, M
+                M_ij = Fi.T @ S[j]
 
-        # The following functions are part of a (rather messy) workaround for computing
-        # a while loop using a for loop with fixed number of iterations.
-        def inner_fn(carry: CarryInnerFn, k: jtp.Int) -> tuple[CarryInnerFn, None]:
-            def compute_inner(carry: CarryInnerFn) -> tuple[CarryInnerFn, None]:
-                j, _, _ = carry
-                out = jax.lax.cond(
-                    pred=(λ[j] > 0),
-                    true_fun=while_loop_body,
-                    false_fun=lambda carry: carry,
-                    operand=carry,
-                )
-                return out, None
+                jj = j - 1
+                M = M.at[ii + 6, jj + 6].set(M_ij.squeeze())
+                M = M.at[jj + 6, ii + 6].set(M_ij.squeeze())
+
+                return j, Fi, M
 
             j, _, _ = carry
-            return jax.lax.cond(
-                pred=(k == j),
-                true_fun=compute_inner,
-                false_fun=lambda carry: (carry, None),
+
+            j, Fi, M = jax.lax.cond(
+                pred=jnp.logical_and(i == λ[j], λ[j] > 0),
+                true_fun=compute,
+                false_fun=lambda carry: carry,
                 operand=carry,
             )
 
+            return (j, Fi, M), None
+
         (j, Fi, M), _ = (
             jax.lax.scan(
-                f=inner_fn,
-                init=carry_inner_fn,
+                f=fake_while_loop,
+                init=fake_while_carry,
                 xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
             )
             if model.number_of_links() > 1

{jaxsim-0.4.3.dev88.dist-info → jaxsim-0.4.3.dev90.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev88
+Version: 0.4.3.dev90
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev88.dist-info → jaxsim-0.4.3.dev90.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=B5nbhwRMLkROhGufRX4ovaMmwptyxOtLB8YwoD_nlo8,426
+jaxsim/_version.py,sha256=fuaYWGJzXhRRB1_d7MJOsOZB0tLvj9VmOhi-3TpDllA,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -47,7 +47,7 @@ jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5
 jaxsim/rbda/__init__.py,sha256=H7DhXpxkPOi9lpUvg31IMHFfRafke1UoJLc5GQIdyhA,387
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
 jaxsim/rbda/collidable_points.py,sha256=Rmf1DhflhOTYh9mDalv0agS0CGSbmfoOybwP2KzKuJ0,4883
-jaxsim/rbda/crba.py,sha256=zJSiHKRvNU98z2tT9prrWR4VU9wIZQWFwEut7mua6as,5044
+jaxsim/rbda/crba.py,sha256=bXkXESnVbv-lxhU1Y_i0rViEcQA4z2t2_jHwdVj5CBo,5049
 jaxsim/rbda/forward_kinematics.py,sha256=2GmEoWsrioVl_SAbKRKfhOLz57pY4aR81PKRdulqStA,3458
 jaxsim/rbda/jacobian.py,sha256=p0EV_8cLzLVV-93VKznT7VPuRj8W7h7rQWkPlWJXfCA,11023
 jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
@@ -63,8 +63,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev88.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev88.dist-info/METADATA,sha256=Hm6Vv2DzMpviXdgHoSsUUuGWJAyYMKngkYXEVMHEh7U,17276
-jaxsim-0.4.3.dev88.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev88.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev88.dist-info/RECORD,,
+jaxsim-0.4.3.dev90.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev90.dist-info/METADATA,sha256=Xtzcvy1beDYPhSFFx_xtt-5fgvGj8c-eeS-_Jv45xfE,17276
+jaxsim-0.4.3.dev90.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev90.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev90.dist-info/RECORD,,