jaxsim 0.4.3.dev231__py3-none-any.whl → 0.4.3.dev242__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
jaxsim/_version.py CHANGED
@@ -12,5 +12,5 @@ __version__: str
12
12
  __version_tuple__: VERSION_TUPLE
13
13
  version_tuple: VERSION_TUPLE
14
14
 
15
- __version__ = version = '0.4.3.dev231'
16
- __version_tuple__ = version_tuple = (0, 4, 3, 'dev231')
15
+ __version__ = version = '0.4.3.dev242'
16
+ __version_tuple__ = version_tuple = (0, 4, 3, 'dev242')
jaxsim/api/contact.py CHANGED
@@ -98,6 +98,7 @@ def collidable_point_forces(
98
98
  data: js.data.JaxSimModelData,
99
99
  link_forces: jtp.MatrixLike | None = None,
100
100
  joint_force_references: jtp.VectorLike | None = None,
101
+ **kwargs,
101
102
  ) -> jtp.Matrix:
102
103
  """
103
104
  Compute the 6D forces applied to each collidable point.
@@ -110,6 +111,7 @@ def collidable_point_forces(
110
111
  representation of data.
111
112
  joint_force_references:
112
113
  The joint force references to apply to the joints.
114
+ kwargs: Additional keyword arguments to pass to the active contact model.
113
115
 
114
116
  Returns:
115
117
  The 6D forces applied to each collidable point expressed in the frame
@@ -121,6 +123,7 @@ def collidable_point_forces(
121
123
  data=data,
122
124
  link_forces=link_forces,
123
125
  joint_force_references=joint_force_references,
126
+ **kwargs,
124
127
  )
125
128
 
126
129
  return f_Ci
@@ -132,7 +135,8 @@ def collidable_point_dynamics(
132
135
  data: js.data.JaxSimModelData,
133
136
  link_forces: jtp.MatrixLike | None = None,
134
137
  joint_force_references: jtp.VectorLike | None = None,
135
- ) -> tuple[jtp.Matrix, dict[str, jtp.Array]]:
138
+ **kwargs,
139
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
136
140
  r"""
137
141
  Compute the 6D force applied to each collidable point.
138
142
 
@@ -144,6 +148,7 @@ def collidable_point_dynamics(
144
148
  representation of data.
145
149
  joint_force_references:
146
150
  The joint force references to apply to the joints.
151
+ kwargs: Additional keyword arguments to pass to the active contact model.
147
152
 
148
153
  Returns:
149
154
  The 6D force applied to each collidable point and additional data based
@@ -158,86 +163,46 @@ def collidable_point_dynamics(
158
163
  Instead, the 6D forces are returned in the active representation.
159
164
  """
160
165
 
161
- # Build the soft contact model.
166
+ # Build the common kw arguments to pass to the computation of the contact forces.
167
+ common_kwargs = dict(
168
+ link_forces=link_forces,
169
+ joint_force_references=joint_force_references,
170
+ )
171
+
172
+ # Build the additional kwargs to pass to the computation of the contact forces.
162
173
  match model.contact_model:
163
174
 
164
175
  case contacts.SoftContacts():
165
- assert isinstance(model.contact_model, contacts.SoftContacts)
166
176
 
167
- # Compute the 6D force expressed in the inertial frame and applied to each
168
- # collidable point, and the corresponding material deformation rate.
169
- # Note that the material deformation rate is always returned in the mixed frame
170
- # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
171
- W_f_Ci, (CW_ṁ,) = model.contact_model.compute_contact_forces(
172
- model=model, data=data
173
- )
174
-
175
- # Create the dictionary of auxiliary data.
176
- # This contact model considers the material deformation as additional state
177
- # of the ODE system. We need to pass its dynamics to the integrator.
178
- aux_data = dict(m_dot=CW_ṁ)
177
+ kwargs_contact_model = {}
179
178
 
180
179
  case contacts.RigidContacts():
181
- assert isinstance(model.contact_model, contacts.RigidContacts)
182
180
 
183
- # Compute the 6D force expressed in the inertial frame and applied to each
184
- # collidable point.
185
- W_f_Ci, _ = model.contact_model.compute_contact_forces(
186
- model=model,
187
- data=data,
188
- link_forces=link_forces,
189
- joint_force_references=joint_force_references,
190
- )
191
-
192
- aux_data = dict()
181
+ kwargs_contact_model = common_kwargs | kwargs
193
182
 
194
183
  case contacts.RelaxedRigidContacts():
195
- assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
196
184
 
197
- # Compute the 6D force expressed in the inertial frame and applied to each
198
- # collidable point.
199
- W_f_Ci, _ = model.contact_model.compute_contact_forces(
200
- model=model,
201
- data=data,
202
- link_forces=link_forces,
203
- joint_force_references=joint_force_references,
204
- )
205
-
206
- aux_data = dict()
185
+ kwargs_contact_model = common_kwargs | kwargs
207
186
 
208
187
  case contacts.ViscoElasticContacts():
209
- assert isinstance(model.contact_model, contacts.ViscoElasticContacts)
210
188
 
211
- # It is not yet clear how to pass the time step to this stage.
212
- # A possibility is to restrict the integrator to only forward Euler
213
- # and store the Δt inside the model.
214
- module = jaxsim.rbda.contacts.visco_elastic.step.__module__
215
- name = jaxsim.rbda.contacts.visco_elastic.step.__name__
216
- msg = "You need to use the custom '{}.{}' function with this contact model."
217
- jaxsim.exceptions.raise_runtime_error_if(
218
- condition=True, msg=msg.format(module, name)
219
- )
220
-
221
- # Compute the 6D force expressed in the inertial frame and applied to each
222
- # collidable point.
223
- W_f_Ci, (W_f̿_Ci, m_tf) = model.contact_model.compute_contact_forces(
224
- model=model,
225
- data=data,
226
- dt=None, # TODO
227
- link_forces=link_forces,
228
- joint_force_references=joint_force_references,
229
- )
230
-
231
- aux_data = dict(W_f_avg2_C=W_f̿_Ci, m_tf=m_tf)
189
+ kwargs_contact_model = common_kwargs | dict(dt=model.time_step) | kwargs
232
190
 
233
191
  case _:
234
- raise ValueError(f"Invalid contact model {model.contact_model}")
192
+ raise ValueError(f"Invalid contact model: {model.contact_model}")
193
+
194
+ # Compute the contact forces with the active contact model.
195
+ W_f_C, aux_data = model.contact_model.compute_contact_forces(
196
+ model=model,
197
+ data=data,
198
+ **kwargs_contact_model,
199
+ )
235
200
 
236
201
  # Compute the transforms of the implicit frames `C[L] = (W_p_C, [L])`
237
202
  # associated to each collidable point.
238
203
  # In inertial-fixed representation, the computation of these transforms
239
204
  # is not necessary and the conversion below becomes a no-op.
240
- W_H_Ci = (
205
+ W_H_C = (
241
206
  js.contact.transforms(model=model, data=data)
242
207
  if data.velocity_representation is not VelRepr.Inertial
243
208
  else jnp.zeros(
@@ -253,7 +218,7 @@ def collidable_point_dynamics(
253
218
  transform=W_H_C,
254
219
  is_force=True,
255
220
  )
256
- )(W_f_Ci, W_H_Ci)
221
+ )(W_f_C, W_H_C)
257
222
 
258
223
  return f_Ci, aux_data
259
224
 
@@ -392,11 +357,13 @@ def estimate_good_contact_parameters(
392
357
  max_penetration=max_δ,
393
358
  number_of_active_collidable_points_steady_state=nc,
394
359
  damping_ratio=damping_ratio,
395
- **dict(
396
- p=model.contact_model.parameters.p,
397
- q=model.contact_model.parameters.q,
398
- )
399
- | kwargs,
360
+ **(
361
+ dict(
362
+ p=model.contact_model.parameters.p,
363
+ q=model.contact_model.parameters.q,
364
+ )
365
+ | kwargs
366
+ ),
400
367
  )
401
368
 
402
369
  case contacts.ViscoElasticContacts():
@@ -410,11 +377,13 @@ def estimate_good_contact_parameters(
410
377
  max_penetration=max_δ,
411
378
  number_of_active_collidable_points_steady_state=nc,
412
379
  damping_ratio=damping_ratio,
413
- **dict(
414
- p=model.contact_model.parameters.p,
415
- q=model.contact_model.parameters.q,
416
- )
417
- | kwargs,
380
+ **(
381
+ dict(
382
+ p=model.contact_model.parameters.p,
383
+ q=model.contact_model.parameters.q,
384
+ )
385
+ | kwargs
386
+ ),
418
387
  )
419
388
  )
420
389
 
@@ -427,11 +396,13 @@ def estimate_good_contact_parameters(
427
396
 
428
397
  parameters = contacts.RigidContactsParams.build(
429
398
  mu=static_friction_coefficient,
430
- **dict(
431
- K=K,
432
- D=2 * jnp.sqrt(K),
433
- )
434
- | kwargs,
399
+ **(
400
+ dict(
401
+ K=K,
402
+ D=2 * jnp.sqrt(K),
403
+ )
404
+ | kwargs
405
+ ),
435
406
  )
436
407
 
437
408
  case contacts.RelaxedRigidContacts():
jaxsim/api/model.py CHANGED
@@ -1770,8 +1770,10 @@ def link_bias_accelerations(
1770
1770
  def link_contact_forces(
1771
1771
  model: js.model.JaxSimModel,
1772
1772
  data: js.data.JaxSimModelData,
1773
+ *,
1773
1774
  link_forces: jtp.MatrixLike | None = None,
1774
1775
  joint_force_references: jtp.VectorLike | None = None,
1776
+ **kwargs,
1775
1777
  ) -> jtp.Matrix:
1776
1778
  """
1777
1779
  Compute the 6D contact forces of all links of the model.
@@ -1784,6 +1786,7 @@ def link_contact_forces(
1784
1786
  representation of data.
1785
1787
  joint_force_references:
1786
1788
  The joint force references to apply to the joints.
1789
+ kwargs: Additional keyword arguments to pass to the active contact model..
1787
1790
 
1788
1791
  Returns:
1789
1792
  A `(nL, 6)` array containing the stacked 6D contact forces of the links,
@@ -1820,47 +1823,16 @@ def link_contact_forces(
1820
1823
  joint_force_references=joint_force_references,
1821
1824
  )
1822
1825
 
1823
- # Compute the 6D forces applied to each collidable point expressed in the
1824
- # inertial frame.
1825
- with (
1826
- data.switch_velocity_representation(VelRepr.Inertial),
1827
- input_references.switch_velocity_representation(VelRepr.Inertial),
1828
- ):
1829
- W_f_C = js.contact.collidable_point_forces(
1830
- model=model,
1831
- data=data,
1832
- link_forces=input_references.link_forces(),
1833
- joint_force_references=input_references.joint_force_references(),
1834
- )
1835
-
1836
- # Construct the vector defining the parent link index of each collidable point.
1837
- # We use this vector to sum the 6D forces of all collidable points rigidly
1838
- # attached to the same link.
1839
- parent_link_index_of_collidable_points = jnp.array(
1840
- model.kin_dyn_parameters.contact_parameters.body, dtype=int
1841
- )
1842
-
1843
- # Create the mask that associate each collidable point to their parent link.
1844
- # We use this mask to sum the collidable points to the right link.
1845
- mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
1846
- model.number_of_links()
1847
- )
1848
-
1849
- # Sum the forces of all collidable points rigidly attached to a body.
1850
- # Since the contact forces W_f_C are expressed in the world frame,
1851
- # we don't need any coordinate transformation.
1852
- W_f_L = mask.T @ W_f_C
1853
-
1854
- # Create a references object to store the link forces.
1855
- references = js.references.JaxSimModelReferences.build(
1856
- model=model, link_forces=W_f_L, velocity_representation=VelRepr.Inertial
1826
+ # Compute the 6D forces applied to the links equivalent to the forces applied
1827
+ # to the frames associated to the collidable points.
1828
+ f_L, _ = model.contact_model.compute_link_contact_forces(
1829
+ model=model,
1830
+ data=data,
1831
+ link_forces=input_references.link_forces(model=model, data=data),
1832
+ joint_force_references=input_references.joint_force_references(),
1833
+ **kwargs,
1857
1834
  )
1858
1835
 
1859
- # Use the references object to convert the link forces to the velocity
1860
- # representation of data.
1861
- with references.switch_velocity_representation(data.velocity_representation):
1862
- f_L = references.link_forces(model=model, data=data)
1863
-
1864
1836
  return f_L
1865
1837
 
1866
1838
 
@@ -1967,6 +1939,11 @@ def step(
1967
1939
  Returns:
1968
1940
  A tuple containing the new data of the model
1969
1941
  and the new state of the integrator.
1942
+
1943
+ Note:
1944
+ In order to reduce the occurrences of frame conversions performed internally,
1945
+ it is recommended to use inertial-fixed velocity representation. This can be
1946
+ particularly useful for automatically differentiated logic.
1970
1947
  """
1971
1948
 
1972
1949
  # Extract the integrator kwargs.
@@ -1976,15 +1953,61 @@ def step(
1976
1953
  integrator_kwargs = kwargs.pop("integrator_kwargs", {})
1977
1954
  integrator_kwargs = kwargs | integrator_kwargs
1978
1955
 
1979
- integrator_state = integrator_state if integrator_state is not None else dict()
1956
+ # Initialize the integrator state.
1957
+ integrator_state_t0 = integrator_state if integrator_state is not None else dict()
1980
1958
 
1981
1959
  # Initialize the time-related variables.
1982
1960
  state_t0 = data.state
1983
1961
  t0 = jnp.array(t0, dtype=float)
1984
1962
  dt = jnp.array(dt if dt is not None else model.time_step).astype(float)
1985
1963
 
1986
- # Rename the integrator state.
1987
- integrator_state_t0 = integrator_state
1964
+ # The visco-elastic contacts operate at best with their own integrator.
1965
+ # They can be used with Euler-like integrators, paying the price of ignoring
1966
+ # some of the benefits of continuous-time integration on the system position.
1967
+ # Furthermore, the requirement to know the Δt used by the integrator is not
1968
+ # compatible with high-order integrators, that use advanced RK stages to evaluate
1969
+ # the dynamics at intermediate times.
1970
+ module = jaxsim.rbda.contacts.visco_elastic.step.__module__
1971
+ name = jaxsim.rbda.contacts.visco_elastic.step.__name__
1972
+ msg = "You need to use the custom '{}.{}' function with this contact model."
1973
+ jaxsim.exceptions.raise_runtime_error_if(
1974
+ condition=(
1975
+ isinstance(model.contact_model, jaxsim.rbda.contacts.ViscoElasticContacts)
1976
+ & (
1977
+ ~jnp.allclose(dt, model.time_step)
1978
+ | ~isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
1979
+ )
1980
+ ),
1981
+ msg=msg.format(module, name),
1982
+ )
1983
+
1984
+ # =================
1985
+ # Phase 1: pre-step
1986
+ # =================
1987
+
1988
+ # TODO: some contact models here may want to perform a dynamic filtering of
1989
+ # the enabled collidable points.
1990
+
1991
+ # Build the references object.
1992
+ # We assume that the link forces are expressed in the frame corresponding to the
1993
+ # velocity representation of the data.
1994
+ references = js.references.JaxSimModelReferences.build(
1995
+ model=model,
1996
+ data=data,
1997
+ velocity_representation=data.velocity_representation,
1998
+ link_forces=link_forces,
1999
+ joint_force_references=joint_force_references,
2000
+ )
2001
+
2002
+ # =============
2003
+ # Phase 2: step
2004
+ # =============
2005
+
2006
+ # Prepare the references to pass.
2007
+ with references.switch_velocity_representation(data.velocity_representation):
2008
+
2009
+ f_L = references.link_forces(model=model, data=data)
2010
+ τ_references = references.joint_force_references(model=model)
1988
2011
 
1989
2012
  # Step the dynamics forward.
1990
2013
  state_tf, integrator_state_tf = integrator.step(
@@ -1994,7 +2017,7 @@ def step(
1994
2017
  params=integrator_state_t0,
1995
2018
  # Always inject the current (model, data) pair into the system dynamics
1996
2019
  # considered by the integrator, and include the input variables represented
1997
- # by the pair (joint_force_references, link_forces).
2020
+ # by the pair (f_L, τ_references).
1998
2021
  # Note that the wrapper of the system dynamics will override (state_x0, t0)
1999
2022
  # inside the passed data even if it is not strictly needed. This logic is
2000
2023
  # necessary to re-use the jit-compiled step function of compatible pytrees
@@ -2003,8 +2026,8 @@ def step(
2003
2026
  dict(
2004
2027
  model=model,
2005
2028
  data=data,
2006
- joint_force_references=joint_force_references,
2007
- link_forces=link_forces,
2029
+ link_forces=f_L,
2030
+ joint_force_references=τ_references,
2008
2031
  )
2009
2032
  | integrator_kwargs
2010
2033
  ),
@@ -2013,6 +2036,10 @@ def step(
2013
2036
  # Store the new state of the model.
2014
2037
  data_tf = data.replace(state=state_tf)
2015
2038
 
2039
+ # ==================
2040
+ # Phase 3: post-step
2041
+ # ==================
2042
+
2016
2043
  # Post process the simulation state, if needed.
2017
2044
  match model.contact_model:
2018
2045
 
@@ -2040,17 +2067,18 @@ def step(
2040
2067
  msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
2041
2068
  )
2042
2069
 
2070
+ W_p_C = js.contact.collidable_point_positions(model, data_tf)
2071
+
2072
+ # Compute the penetration depth of the collidable points.
2073
+ δ, *_ = jax.vmap(
2074
+ jaxsim.rbda.contacts.common.compute_penetration_data,
2075
+ in_axes=(0, 0, None),
2076
+ )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
2077
+
2043
2078
  with data_tf.switch_velocity_representation(VelRepr.Mixed):
2044
2079
 
2045
2080
  J_WC = js.contact.jacobian(model, data_tf)
2046
2081
  M = js.model.free_floating_mass_matrix(model, data_tf)
2047
- W_p_C = js.contact.collidable_point_positions(model, data_tf)
2048
-
2049
- # Compute the penetration depth of the collidable points.
2050
- δ, *_ = jax.vmap(
2051
- jaxsim.rbda.contacts.common.compute_penetration_data,
2052
- in_axes=(0, 0, None),
2053
- )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
2054
2082
 
2055
2083
  # Compute the impact velocity.
2056
2084
  # It may be discontinuous in case new contacts are made.
@@ -2063,13 +2091,13 @@ def step(
2063
2091
  )
2064
2092
  )
2065
2093
 
2066
- # Reset the generalized velocity.
2067
- data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
2068
- data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
2094
+ # Reset the generalized velocity.
2095
+ data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
2096
+ data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
2069
2097
 
2070
- # Restore the input velocity representation.
2071
- data_tf = data_tf.replace(
2072
- velocity_representation=data.velocity_representation, validate=False
2073
- )
2098
+ # Restore the input velocity representation.
2099
+ data_tf = data_tf.replace(
2100
+ velocity_representation=data.velocity_representation, validate=False
2101
+ )
2074
2102
 
2075
2103
  return data_tf, integrator_state_tf
jaxsim/api/ode.py CHANGED
@@ -131,7 +131,7 @@ def system_velocity_dynamics(
131
131
 
132
132
  # Initialize the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
133
133
  # with the terrain.
134
- W_f_Li_terrain = jnp.zeros_like(O_f_L).astype(float)
134
+ W_f_L_terrain = jnp.zeros_like(O_f_L).astype(float)
135
135
 
136
136
  # Initialize a dictionary of auxiliary data.
137
137
  # This dictionary is used to store additional data computed by the contact model.
@@ -139,66 +139,59 @@ def system_velocity_dynamics(
139
139
 
140
140
  if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
141
141
 
142
- # Note: the following code should be kept in sync with the function
143
- # `jaxsim.api.model.link_contact_forces`. We cannot merge them since
144
- # here we need to get also aux_data.
145
-
146
- # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
147
- # along with contact-specific auxiliary states.
148
142
  with (
149
143
  data.switch_velocity_representation(VelRepr.Inertial),
150
144
  references.switch_velocity_representation(VelRepr.Inertial),
151
145
  ):
152
- W_f_Ci, aux_data = js.contact.collidable_point_dynamics(
146
+
147
+ # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
148
+ # along with contact-specific auxiliary states.
149
+ W_f_C, aux_data = js.contact.collidable_point_dynamics(
153
150
  model=model,
154
151
  data=data,
155
152
  link_forces=references.link_forces(model=model, data=data),
156
153
  joint_force_references=references.joint_force_references(model=model),
157
154
  )
158
155
 
159
- # Construct the vector defining the parent link index of each collidable point.
160
- # We use this vector to sum the 6D forces of all collidable points rigidly
161
- # attached to the same link.
162
- parent_link_index_of_collidable_points = jnp.array(
163
- model.kin_dyn_parameters.contact_parameters.body, dtype=int
164
- )
165
-
166
- # Sum the forces of all collidable points rigidly attached to a body.
167
- # Since the contact forces W_f_Ci are expressed in the world frame,
168
- # we don't need any coordinate transformation.
169
- mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
170
- model.number_of_links()
171
- )
172
-
173
- W_f_Li_terrain = mask.T @ W_f_Ci
156
+ # Compute the 6D forces applied to the links equivalent to the forces applied
157
+ # to the frames associated to the collidable points.
158
+ W_f_L_terrain = model.contact_model.link_forces_from_contact_forces(
159
+ model=model,
160
+ data=data,
161
+ contact_forces=W_f_C,
162
+ )
174
163
 
175
164
  # ===========================
176
165
  # Compute system acceleration
177
166
  # ===========================
178
167
 
179
- # Compute the total link forces
168
+ # Compute the total link forces.
180
169
  with (
181
170
  data.switch_velocity_representation(VelRepr.Inertial),
182
171
  references.switch_velocity_representation(VelRepr.Inertial),
183
172
  ):
173
+
174
+ # Sum the contact forces just computed with the link forces applied by the user.
184
175
  references = references.apply_link_forces(
185
176
  model=model,
186
177
  data=data,
187
- forces=W_f_Li_terrain,
178
+ forces=W_f_L_terrain,
188
179
  additive=True,
189
180
  )
190
181
 
191
- # Get the link forces in inertial representation
182
+ # Get the link forces in inertial-fixed representation.
192
183
  f_L_total = references.link_forces(model=model, data=data)
193
184
 
194
- v̇_WB, = system_acceleration(
185
+ # Compute the system acceleration in inertial-fixed representation.
186
+ # This representation is useful for integration purpose.
187
+ W_v̇_WB, s̈ = system_acceleration(
195
188
  model=model,
196
189
  data=data,
197
190
  joint_force_references=joint_force_references,
198
191
  link_forces=f_L_total,
199
192
  )
200
193
 
201
- return v̇_WB, s̈, aux_data
194
+ return W_v̇_WB, s̈, aux_data
202
195
 
203
196
 
204
197
  def system_acceleration(
@@ -390,17 +383,15 @@ def system_dynamics(
390
383
 
391
384
  case contacts.ViscoElasticContacts():
392
385
 
393
- extended_ode_state["contacts_state"] = {
394
- "tangential_deformation": jnp.zeros_like(
395
- data.state.extended["tangential_deformation"]
396
- )
397
- }
386
+ extended_ode_state["tangential_deformation"] = jnp.zeros_like(
387
+ data.state.extended["tangential_deformation"]
388
+ )
398
389
 
399
390
  case contacts.RigidContacts() | contacts.RelaxedRigidContacts():
400
391
  pass
401
392
 
402
393
  case _:
403
- raise ValueError(f"Invalid contact model {model.contact_model}")
394
+ raise ValueError(f"Invalid contact model: {model.contact_model}")
404
395
 
405
396
  # Extract the velocities.
406
397
  W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
@@ -2,7 +2,6 @@ from __future__ import annotations
2
2
 
3
3
  import abc
4
4
  import functools
5
- from typing import Any
6
5
 
7
6
  import jax
8
7
  import jax.numpy as jnp
@@ -10,6 +9,7 @@ import jax.numpy as jnp
10
9
  import jaxsim.api as js
11
10
  import jaxsim.terrain
12
11
  import jaxsim.typing as jtp
12
+ from jaxsim.api.common import ModelDataWithVelocityRepresentation
13
13
  from jaxsim.utils import JaxsimDataclass
14
14
 
15
15
  try:
@@ -131,7 +131,7 @@ class ContactModel(JaxsimDataclass):
131
131
  model: js.model.JaxSimModel,
132
132
  data: js.data.JaxSimModelData,
133
133
  **kwargs,
134
- ) -> tuple[jtp.Matrix, tuple[Any, ...]]:
134
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
135
135
  """
136
136
  Compute the contact forces.
137
137
 
@@ -142,11 +142,145 @@ class ContactModel(JaxsimDataclass):
142
142
  Returns:
143
143
  A tuple containing as first element the computed 6D contact force applied to
144
144
  the contact points and expressed in the world frame, and as second element
145
- a tuple of optional additional information.
145
+ a dictionary of optional additional information.
146
146
  """
147
147
 
148
148
  pass
149
149
 
150
+ def compute_link_contact_forces(
151
+ self,
152
+ model: js.model.JaxSimModel,
153
+ data: js.data.JaxSimModelData,
154
+ **kwargs,
155
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
156
+ """
157
+ Compute the link contact forces.
158
+
159
+ Args:
160
+ model: The robot model considered by the contact model.
161
+ data: The data of the considered model.
162
+
163
+ Returns:
164
+ A tuple containing as first element the 6D contact force applied to the
165
+ links and expressed in the frame of the velocity representation of data,
166
+ and as second element a dictionary of optional additional information.
167
+ """
168
+
169
+ # Compute the contact forces expressed in the inertial frame.
170
+ # This function, contrarily to `compute_contact_forces`, already handles how
171
+ # the optional kwargs should be passed to the specific contact models.
172
+ W_f_C, aux_dict = js.contact.collidable_point_dynamics(
173
+ model=model, data=data, **kwargs
174
+ )
175
+
176
+ # Compute the 6D forces applied to the links equivalent to the forces applied
177
+ # to the frames associated to the collidable points.
178
+ with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
179
+
180
+ W_f_L = self.link_forces_from_contact_forces(
181
+ model=model, data=data, contact_forces=W_f_C
182
+ )
183
+
184
+ # Store the link forces in the references object for easy conversion.
185
+ references = js.references.JaxSimModelReferences.build(
186
+ model=model,
187
+ data=data,
188
+ link_forces=W_f_L,
189
+ velocity_representation=jaxsim.VelRepr.Inertial,
190
+ )
191
+
192
+ # Convert the link forces to the frame corresponding to the velocity
193
+ # representation of data.
194
+ with references.switch_velocity_representation(data.velocity_representation):
195
+ f_L = references.link_forces(model=model, data=data)
196
+
197
+ return f_L, aux_dict
198
+
199
+ @staticmethod
200
+ def link_forces_from_contact_forces(
201
+ model: js.model.JaxSimModel,
202
+ data: js.data.JaxSimModelData,
203
+ *,
204
+ contact_forces: jtp.MatrixLike,
205
+ ) -> jtp.Matrix:
206
+ """
207
+ Compute the link forces from the contact forces.
208
+
209
+ Args:
210
+ model: The robot model considered by the contact model.
211
+ data: The data of the considered model.
212
+ contact_forces: The contact forces computed by the contact model.
213
+
214
+ Returns:
215
+ The 6D contact forces applied to the links and expressed in the frame of
216
+ the velocity representation of data.
217
+ """
218
+
219
+ # Convert the contact forces to a JAX array.
220
+ f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
221
+
222
+ # Get the pose of the enabled collidable points.
223
+ W_H_C = js.contact.transforms(model=model, data=data)
224
+
225
+ # Convert the contact forces to inertial-fixed representation.
226
+ W_f_C = jax.vmap(
227
+ lambda f_C, W_H_C: (
228
+ ModelDataWithVelocityRepresentation.other_representation_to_inertial(
229
+ array=f_C,
230
+ other_representation=data.velocity_representation,
231
+ transform=W_H_C,
232
+ is_force=True,
233
+ )
234
+ )
235
+ )(f_C, W_H_C)
236
+
237
+ # Get the object storing the contact parameters of the model.
238
+ contact_parameters = model.kin_dyn_parameters.contact_parameters
239
+
240
+ # Extract the indices corresponding to the enabled collidable points.
241
+ indices_of_enabled_collidable_points = (
242
+ contact_parameters.indices_of_enabled_collidable_points
243
+ )
244
+
245
+ # Construct the vector defining the parent link index of each collidable point.
246
+ # We use this vector to sum the 6D forces of all collidable points rigidly
247
+ # attached to the same link.
248
+ parent_link_index_of_collidable_points = jnp.array(
249
+ contact_parameters.body, dtype=int
250
+ )[indices_of_enabled_collidable_points]
251
+
252
+ # Create the mask that associate each collidable point to their parent link.
253
+ # We use this mask to sum the collidable points to the right link.
254
+ mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
255
+ model.number_of_links()
256
+ )
257
+
258
+ # Sum the forces of all collidable points rigidly attached to a body.
259
+ # Since the contact forces W_f_C are expressed in the world frame,
260
+ # we don't need any coordinate transformation.
261
+ W_f_L = mask.T @ W_f_C
262
+
263
+ # Compute the link transforms.
264
+ W_H_L = (
265
+ js.model.forward_kinematics(model=model, data=data)
266
+ if data.velocity_representation is not jaxsim.VelRepr.Inertial
267
+ else jnp.zeros(shape=(model.number_of_links(), 4, 4))
268
+ )
269
+
270
+ # Convert the inertial-fixed link forces to the velocity representation of data.
271
+ f_L = jax.vmap(
272
+ lambda W_f_L, W_H_L: (
273
+ ModelDataWithVelocityRepresentation.inertial_to_other_representation(
274
+ array=W_f_L,
275
+ other_representation=data.velocity_representation,
276
+ transform=W_H_L,
277
+ is_force=True,
278
+ )
279
+ )
280
+ )(W_f_L, W_H_L)
281
+
282
+ return f_L
283
+
150
284
  @classmethod
151
285
  def zero_state_variables(cls, model: js.model.JaxSimModel) -> dict[str, jtp.Array]:
152
286
  """
@@ -120,19 +120,44 @@ class RelaxedRigidContactsParams(common.ContactsParams):
120
120
 
121
121
  return cls(
122
122
  time_constant=jnp.array(
123
- time_constant or default("time_constant"), dtype=float
123
+ (
124
+ time_constant
125
+ if time_constant is not None
126
+ else default("time_constant")
127
+ ),
128
+ dtype=float,
124
129
  ),
125
130
  damping_coefficient=jnp.array(
126
- damping_coefficient or default("damping_coefficient"), dtype=float
131
+ (
132
+ damping_coefficient
133
+ if damping_coefficient is not None
134
+ else default("damping_coefficient")
135
+ ),
136
+ dtype=float,
137
+ ),
138
+ d_min=jnp.array(
139
+ d_min if d_min is not None else default("d_min"), dtype=float
140
+ ),
141
+ d_max=jnp.array(
142
+ d_max if d_max is not None else default("d_max"), dtype=float
143
+ ),
144
+ width=jnp.array(
145
+ width if width is not None else default("width"), dtype=float
146
+ ),
147
+ midpoint=jnp.array(
148
+ midpoint if midpoint is not None else default("midpoint"), dtype=float
127
149
  ),
128
- d_min=jnp.array(d_min or default("d_min"), dtype=float),
129
- d_max=jnp.array(d_max or default("d_max"), dtype=float),
130
- width=jnp.array(width or default("width"), dtype=float),
131
- midpoint=jnp.array(midpoint or default("midpoint"), dtype=float),
132
- power=jnp.array(power or default("power"), dtype=float),
133
- stiffness=jnp.array(stiffness or default("stiffness"), dtype=float),
134
- damping=jnp.array(damping or default("damping"), dtype=float),
135
- mu=jnp.array(mu or default("mu"), dtype=float),
150
+ power=jnp.array(
151
+ power if power is not None else default("power"), dtype=float
152
+ ),
153
+ stiffness=jnp.array(
154
+ stiffness if stiffness is not None else default("stiffness"),
155
+ dtype=float,
156
+ ),
157
+ damping=jnp.array(
158
+ damping if damping is not None else default("damping"), dtype=float
159
+ ),
160
+ mu=jnp.array(mu if mu is not None else default("mu"), dtype=float),
136
161
  )
137
162
 
138
163
  def valid(self) -> jtp.BoolLike:
@@ -210,7 +235,9 @@ class RelaxedRigidContacts(common.ContactModel):
210
235
 
211
236
  # Create the solver options to set by combining the default solver options
212
237
  # with the user-provided solver options.
213
- solver_options = default_solver_options | (solver_options or {})
238
+ solver_options = default_solver_options | (
239
+ solver_options if solver_options is not None else {}
240
+ )
214
241
 
215
242
  # Make sure that the solver options are hashable.
216
243
  # We need to check this because the solver options are static.
@@ -223,9 +250,15 @@ class RelaxedRigidContacts(common.ContactModel):
223
250
 
224
251
  return cls(
225
252
  parameters=(
226
- parameters or cls.__dataclass_fields__["parameters"].default_factory()
253
+ parameters
254
+ if parameters is not None
255
+ else cls.__dataclass_fields__["parameters"].default_factory()
256
+ ),
257
+ terrain=(
258
+ terrain
259
+ if terrain is not None
260
+ else cls.__dataclass_fields__["terrain"].default_factory()
227
261
  ),
228
- terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
229
262
  _solver_options_keys=tuple(solver_options.keys()),
230
263
  _solver_options_values=tuple(solver_options.values()),
231
264
  )
@@ -238,7 +271,7 @@ class RelaxedRigidContacts(common.ContactModel):
238
271
  *,
239
272
  link_forces: jtp.MatrixLike | None = None,
240
273
  joint_force_references: jtp.VectorLike | None = None,
241
- ) -> tuple[jtp.Matrix, tuple]:
274
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
242
275
  """
243
276
  Compute the contact forces.
244
277
 
@@ -458,7 +491,7 @@ class RelaxedRigidContacts(common.ContactModel):
458
491
  ),
459
492
  )(CW_fl_C, W_H_C)
460
493
 
461
- return W_f_C, ()
494
+ return W_f_C, {}
462
495
 
463
496
  @staticmethod
464
497
  def _regularizers(
@@ -66,9 +66,17 @@ class RigidContactsParams(ContactsParams):
66
66
  """Create a `RigidContactParams` instance"""
67
67
 
68
68
  return cls(
69
- mu=mu or cls.__dataclass_fields__["mu"].default,
70
- K=K or cls.__dataclass_fields__["K"].default,
71
- D=D or cls.__dataclass_fields__["D"].default,
69
+ mu=jnp.array(
70
+ mu
71
+ if mu is not None
72
+ else cls.__dataclass_fields__["mu"].default_factory()
73
+ ).astype(float),
74
+ K=jnp.array(
75
+ K if K is not None else cls.__dataclass_fields__["K"].default_factory()
76
+ ).astype(float),
77
+ D=jnp.array(
78
+ D if D is not None else cls.__dataclass_fields__["D"].default_factory()
79
+ ).astype(float),
72
80
  )
73
81
 
74
82
  def valid(self) -> jtp.BoolLike:
@@ -147,7 +155,9 @@ class RigidContacts(ContactModel):
147
155
 
148
156
  # Create the solver options to set by combining the default solver options
149
157
  # with the user-provided solver options.
150
- solver_options = default_solver_options | (solver_options or {})
158
+ solver_options = default_solver_options | (
159
+ solver_options if solver_options is not None else {}
160
+ )
151
161
 
152
162
  # Make sure that the solver options are hashable.
153
163
  # We need to check this because the solver options are static.
@@ -160,12 +170,19 @@ class RigidContacts(ContactModel):
160
170
 
161
171
  return cls(
162
172
  parameters=(
163
- parameters or cls.__dataclass_fields__["parameters"].default_factory()
173
+ parameters
174
+ if parameters is not None
175
+ else cls.__dataclass_fields__["parameters"].default_factory()
176
+ ),
177
+ terrain=(
178
+ terrain
179
+ if terrain is not None
180
+ else cls.__dataclass_fields__["terrain"].default_factory()
164
181
  ),
165
- terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
166
182
  regularization_delassus=float(
167
183
  regularization_delassus
168
- or cls.__dataclass_fields__["regularization_delassus"].default
184
+ if regularization_delassus is not None
185
+ else cls.__dataclass_fields__["regularization_delassus"].default
169
186
  ),
170
187
  _solver_options_keys=tuple(solver_options.keys()),
171
188
  _solver_options_values=tuple(solver_options.values()),
@@ -242,7 +259,7 @@ class RigidContacts(ContactModel):
242
259
  *,
243
260
  link_forces: jtp.MatrixLike | None = None,
244
261
  joint_force_references: jtp.VectorLike | None = None,
245
- ) -> tuple[jtp.Matrix, tuple]:
262
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
246
263
  """
247
264
  Compute the contact forces.
248
265
 
@@ -402,7 +419,7 @@ class RigidContacts(ContactModel):
402
419
  ),
403
420
  )(CW_fl_C, W_H_C)
404
421
 
405
- return W_f_C, ()
422
+ return W_f_C, {}
406
423
 
407
424
  @staticmethod
408
425
  def _delassus_matrix(
@@ -237,9 +237,13 @@ class SoftContacts(common.ContactModel):
237
237
  else cls.__dataclass_fields__["parameters"].default_factory()
238
238
  )
239
239
 
240
- return SoftContacts(
240
+ return cls(
241
241
  parameters=parameters,
242
- terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
242
+ terrain=(
243
+ terrain
244
+ if terrain is not None
245
+ else cls.__dataclass_fields__["terrain"].default_factory()
246
+ ),
243
247
  )
244
248
 
245
249
  @classmethod
@@ -423,7 +427,7 @@ class SoftContacts(common.ContactModel):
423
427
  self,
424
428
  model: js.model.JaxSimModel,
425
429
  data: js.data.JaxSimModelData,
426
- ) -> tuple[jtp.Matrix, tuple[jtp.Matrix]]:
430
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
427
431
  """
428
432
  Compute the contact forces.
429
433
 
@@ -433,7 +437,7 @@ class SoftContacts(common.ContactModel):
433
437
 
434
438
  Returns:
435
439
  A tuple containing as first element the computed contact forces, and as
436
- second element the derivative of the material deformation.
440
+ second element a dictionary with derivative of the material deformation.
437
441
  """
438
442
 
439
443
  # Initialize the model and data this contact model is operating on.
@@ -460,4 +464,4 @@ class SoftContacts(common.ContactModel):
460
464
  )
461
465
  )(W_p_C, W_ṗ_C, m)
462
466
 
463
- return W_f, (ṁ,)
467
+ return W_f, dict(m_dot=ṁ)
@@ -13,6 +13,7 @@ import jaxsim.api as js
13
13
  import jaxsim.exceptions
14
14
  import jaxsim.typing as jtp
15
15
  from jaxsim import logging
16
+ from jaxsim.api.common import ModelDataWithVelocityRepresentation
16
17
  from jaxsim.math import StandardGravity
17
18
  from jaxsim.terrain import FlatTerrain, Terrain
18
19
 
@@ -235,11 +236,17 @@ class ViscoElasticContacts(common.ContactModel):
235
236
  else cls.__dataclass_fields__["parameters"].default_factory()
236
237
  )
237
238
 
238
- return ViscoElasticContacts(
239
+ return cls(
239
240
  parameters=parameters,
240
- terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
241
+ terrain=(
242
+ terrain
243
+ if terrain is not None
244
+ else cls.__dataclass_fields__["terrain"].default_factory()
245
+ ),
241
246
  max_squarings=int(
242
- max_squarings or cls.__dataclass_fields__["max_squarings"].default
247
+ max_squarings
248
+ if max_squarings is not None
249
+ else cls.__dataclass_fields__["max_squarings"].default
243
250
  ),
244
251
  )
245
252
 
@@ -266,7 +273,7 @@ class ViscoElasticContacts(common.ContactModel):
266
273
  dt: jtp.FloatLike | None = None,
267
274
  link_forces: jtp.MatrixLike | None = None,
268
275
  joint_force_references: jtp.VectorLike | None = None,
269
- ) -> tuple[jtp.Matrix, tuple[jtp.Matrix, jtp.Matrix]]:
276
+ ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
270
277
  """
271
278
  Compute the contact forces.
272
279
 
@@ -291,7 +298,7 @@ class ViscoElasticContacts(common.ContactModel):
291
298
  Returns:
292
299
  A tuple containing as first element the computed 6D contact force applied to
293
300
  the contact point and expressed in the world frame, and as second element
294
- a tuple of optional additional information.
301
+ a dictionary of optional additional information.
295
302
  """
296
303
 
297
304
  # Initialize the model and data this contact model is operating on.
@@ -315,8 +322,8 @@ class ViscoElasticContacts(common.ContactModel):
315
322
  model=model,
316
323
  data=data,
317
324
  dt=jnp.array(dt).astype(float),
318
- joint_force_references=joint_force_references,
319
325
  link_forces=link_forces,
326
+ joint_force_references=joint_force_references,
320
327
  indices_of_enabled_collidable_points=indices_of_enabled_collidable_points,
321
328
  max_squarings=self.max_squarings,
322
329
  )
@@ -334,11 +341,13 @@ class ViscoElasticContacts(common.ContactModel):
334
341
 
335
342
  # Vmapped transformation from mixed to inertial-fixed representation.
336
343
  compute_forces_inertial_fixed_vmap = jax.vmap(
337
- lambda CW_fl_C, W_H_C: data.other_representation_to_inertial(
338
- array=jnp.zeros(6).at[0:3].set(CW_fl_C),
339
- other_representation=jaxsim.VelRepr.Mixed,
340
- transform=W_H_C,
341
- is_force=True,
344
+ lambda CW_fl_C, W_H_C: (
345
+ ModelDataWithVelocityRepresentation.other_representation_to_inertial(
346
+ array=jnp.zeros(6).at[0:3].set(CW_fl_C),
347
+ other_representation=jaxsim.VelRepr.Mixed,
348
+ transform=W_H_C,
349
+ is_force=True,
350
+ )
342
351
  )
343
352
  )
344
353
 
@@ -347,7 +356,7 @@ class ViscoElasticContacts(common.ContactModel):
347
356
  lambda CW_fl: compute_forces_inertial_fixed_vmap(CW_fl, W_H_C)
348
357
  )(jnp.stack([CW_f̅l, CW_fl̿]))
349
358
 
350
- return W_f̅_C, (W_f̿_C, m_tf)
359
+ return W_f̅_C, dict(W_f_avg2_C=W_f̿_C, m_tf=m_tf)
351
360
 
352
361
  @staticmethod
353
362
  @functools.partial(jax.jit, static_argnames=("max_squarings",))
@@ -407,8 +416,8 @@ class ViscoElasticContacts(common.ContactModel):
407
416
  A, b, A_sc, b_sc = ViscoElasticContacts._contact_points_dynamics(
408
417
  model=model,
409
418
  data=data,
410
- joint_force_references=joint_force_references,
411
419
  link_forces=link_forces,
420
+ joint_force_references=joint_force_references,
412
421
  indices_of_enabled_collidable_points=indices,
413
422
  p_t0=p_t0,
414
423
  v_t0=v_t0,
@@ -657,8 +666,8 @@ class ViscoElasticContacts(common.ContactModel):
657
666
  BW_v̇_free_WB, s̈_free = js.ode.system_acceleration(
658
667
  model=model,
659
668
  data=data,
660
- joint_force_references=references.joint_force_references(model=model),
661
669
  link_forces=references.link_forces(model=model, data=data),
670
+ joint_force_references=references.joint_force_references(model=model),
662
671
  )
663
672
 
664
673
  # Pack the free system acceleration in mixed representation.
@@ -688,7 +697,20 @@ class ViscoElasticContacts(common.ContactModel):
688
697
  parameters: ViscoElasticContactsParams,
689
698
  terrain: Terrain,
690
699
  ) -> tuple[jtp.Matrix, jtp.Vector]:
691
- """"""
700
+ """
701
+ Linearize the Hunt/Crossley contact model at the initial state.
702
+
703
+ Args:
704
+ position: The position of the contact point.
705
+ velocity: The velocity of the contact point.
706
+ tangential_deformation: The tangential deformation of the contact point.
707
+ parameters: The parameters of the contact model.
708
+ terrain: The considered terrain.
709
+
710
+ Returns:
711
+ A tuple containing the `A` matrix and the `b` vector of the linear system
712
+ corresponding to the contact dynamics linearized at the initial state.
713
+ """
692
714
 
693
715
  # Initialize the state at which the model is linearized.
694
716
  p0 = jnp.array(position, dtype=float).squeeze()
@@ -969,58 +991,67 @@ def step(
969
991
  assert isinstance(model.contact_model, ViscoElasticContacts)
970
992
  assert isinstance(data.contacts_params, ViscoElasticContactsParams)
971
993
 
994
+ # Compute the contact forces in inertial-fixed representation.
995
+ # TODO: understand what's wrong in other representations.
996
+ data_inertial_fixed = data.replace(
997
+ velocity_representation=jaxsim.VelRepr.Inertial, validate=False
998
+ )
999
+
1000
+ # Create the references object.
1001
+ references = js.references.JaxSimModelReferences.build(
1002
+ model=model,
1003
+ data=data,
1004
+ link_forces=link_forces,
1005
+ joint_force_references=joint_force_references,
1006
+ velocity_representation=data.velocity_representation,
1007
+ )
1008
+
972
1009
  # Initialize the time step.
973
1010
  dt = dt if dt is not None else model.time_step
974
1011
 
975
1012
  # Compute the contact forces with the exponential integrator.
976
- W_f̅_C, (W_f̿_C, m_tf) = model.contact_model.compute_contact_forces(
1013
+ W_f̅_C, aux_data = model.contact_model.compute_contact_forces(
977
1014
  model=model,
978
- data=data,
1015
+ data=data_inertial_fixed,
979
1016
  dt=jnp.array(dt).astype(float),
980
- link_forces=link_forces,
981
- joint_force_references=joint_force_references,
1017
+ link_forces=references.link_forces(model=model, data=data),
1018
+ joint_force_references=references.joint_force_references(model=model),
982
1019
  )
983
1020
 
1021
+ # Extract the final material deformation and the average of average forces
1022
+ # from the dictionary containing auxiliary data.
1023
+ m_tf = aux_data["m_tf"]
1024
+ W_f̿_C = aux_data["W_f_avg2_C"]
1025
+
984
1026
  # ===============================
985
1027
  # Compute the link contact forces
986
1028
  # ===============================
987
1029
 
988
- # Extract the indices corresponding to the enabled collidable points.
989
- # The visco-elastic contact model computed only their contact forces.
990
- indices_of_enabled_collidable_points = (
991
- model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
992
- )
1030
+ # Get the link contact forces by summing the forces of contact points belonging
1031
+ # to the same link.
1032
+ W_f̅_L, W_f̿_L = jax.vmap(
1033
+ lambda W_f_C: model.contact_model.link_forces_from_contact_forces(
1034
+ model=model, data=data_inertial_fixed, contact_forces=W_f_C
1035
+ )
1036
+ )(jnp.stack([W_f̅_C, W_f̿_C]))
993
1037
 
994
1038
  # Compute the link transforms.
995
- W_H_L = js.model.forward_kinematics(model=model, data=data)
996
-
997
- # Construct the vector defining the parent link index of each collidable point.
998
- # We use this vector to sum the 6D forces of all collidable points rigidly
999
- # attached to the same link.
1000
- parent_link_index_of_collidable_points = jnp.array(
1001
- model.kin_dyn_parameters.contact_parameters.body, dtype=int
1002
- )[indices_of_enabled_collidable_points]
1003
-
1004
- # Create the mask that associate each collidable point to their parent link.
1005
- # We use this mask to sum the collidable points to the right link.
1006
- mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
1007
- model.number_of_links()
1039
+ W_H_L = (
1040
+ js.model.forward_kinematics(model=model, data=data)
1041
+ if data.velocity_representation is not jaxsim.VelRepr.Inertial
1042
+ else jnp.zeros(shape=(model.number_of_links(), 4, 4))
1008
1043
  )
1009
1044
 
1010
- # Sum the forces of all collidable points rigidly attached to a body.
1011
- # Since the contact forces W_f_C are expressed in the world frame,
1012
- # we don't need any coordinate transformation.
1013
- W_f̅_L = mask.T @ W_f̅_C
1014
- W_f̿_L = mask.T @ W_f̿_C
1015
-
1016
- # For integration purpose, we need these average of averages expressed in
1045
+ # For integration purpose, we need the average of average forces expressed in
1017
1046
  # mixed representation.
1018
1047
  LW_f̿_L = jax.vmap(
1019
- lambda W_f_L, W_H_L: data.inertial_to_other_representation(
1020
- array=W_f_L,
1021
- other_representation=jaxsim.VelRepr.Mixed,
1022
- transform=W_H_L,
1023
- is_force=True,
1048
+ lambda W_f_L, W_H_L: (
1049
+ ModelDataWithVelocityRepresentation.inertial_to_other_representation(
1050
+ array=W_f_L,
1051
+ other_representation=jaxsim.VelRepr.Mixed,
1052
+ transform=W_H_L,
1053
+ is_force=True,
1054
+ )
1024
1055
  )
1025
1056
  )(W_f̿_L, W_H_L)
1026
1057
 
@@ -1032,10 +1063,10 @@ def step(
1032
1063
  data_tf: js.data.JaxSimModelData = (
1033
1064
  model.contact_model.integrate_data_with_average_contact_forces(
1034
1065
  model=model,
1035
- data=data,
1066
+ data=data_inertial_fixed,
1036
1067
  dt=dt,
1037
- link_forces=link_forces,
1038
- joint_force_references=joint_force_references,
1068
+ link_forces=references.link_forces(model=model, data=data),
1069
+ joint_force_references=references.joint_force_references(model=model),
1039
1070
  average_link_contact_forces_inertial=W_f̅_L,
1040
1071
  average_of_average_link_contact_forces_mixed=LW_f̿_L,
1041
1072
  )
@@ -1046,10 +1077,21 @@ def step(
1046
1077
  # be much more accurate than the one computed with the discrete soft contacts.
1047
1078
  with data_tf.mutable_context():
1048
1079
 
1080
+ # Extract the indices corresponding to the enabled collidable points.
1081
+ # The visco-elastic contact model computed only their contact forces.
1082
+ indices_of_enabled_collidable_points = (
1083
+ model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
1084
+ )
1085
+
1049
1086
  data_tf.state.extended |= {
1050
1087
  "tangential_deformation": data_tf.state.extended["tangential_deformation"]
1051
1088
  .at[indices_of_enabled_collidable_points]
1052
1089
  .set(m_tf)
1053
1090
  }
1054
1091
 
1092
+ # Restore the original velocity representation.
1093
+ data_tf = data_tf.replace(
1094
+ velocity_representation=data.velocity_representation, validate=False
1095
+ )
1096
+
1055
1097
  return data_tf, {}
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: jaxsim
3
- Version: 0.4.3.dev231
3
+ Version: 0.4.3.dev242
4
4
  Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
5
5
  Author-email: Diego Ferigo <dgferigo@gmail.com>
6
6
  Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>
@@ -1,19 +1,19 @@
1
1
  jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
2
- jaxsim/_version.py,sha256=1YX6XhCYEbtZpJ6_hP41mzRQ3Iq0wjcPnL1r96jz93I,428
2
+ jaxsim/_version.py,sha256=MulyCITmYYLdFRMIlPKDiFIXlIDsYkQ31wXsix-6zyo,428
3
3
  jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
4
4
  jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
5
5
  jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
6
6
  jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
7
7
  jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
8
8
  jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
9
- jaxsim/api/contact.py,sha256=Egc62310ljn5goXlswwJYSB-LyW6M5gmPoT_a3mkd7U,25812
9
+ jaxsim/api/contact.py,sha256=iiQc8vdlJ60Vl-ki4ieP7TW6MJhJfWREYpcVyyEWT0M,24034
10
10
  jaxsim/api/data.py,sha256=gQX6hfEaw0ooJYvpr5f8UvEJwqhtflEK_NHWn9XgTZY,28935
11
11
  jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
12
12
  jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
13
13
  jaxsim/api/kin_dyn_parameters.py,sha256=thJbz9XhpXgom23S6MXX2ugxGoAD-k947ZMAHDisy2w,29620
14
14
  jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
15
- jaxsim/api/model.py,sha256=s2i4obxMjZ_XntJgT0dEV57LCo0GIC7VppUnxsqC1fc,69704
16
- jaxsim/api/ode.py,sha256=J_WuaoPl3ZY-yvTrCQun-rQoIAv_duynSXAGxqx93sg,14211
15
+ jaxsim/api/model.py,sha256=FSK2KvH2oNoGIt67gxQD7ScbClqXCn1Oy5JLtc4QHYg,70585
16
+ jaxsim/api/ode.py,sha256=2KvGT3WW1eWEme4fH-5LlOXdLF6JUP5Z-IGY93ashUc,13815
17
17
  jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
18
18
  jaxsim/api/references.py,sha256=fW77LitZ8DYgT6ZmUInJfm5luBV1mTcqcNRiC_i79og,20862
19
19
  jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
@@ -54,19 +54,19 @@ jaxsim/rbda/jacobian.py,sha256=p0EV_8cLzLVV-93VKznT7VPuRj8W7h7rQWkPlWJXfCA,11023
54
54
  jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
55
55
  jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
56
56
  jaxsim/rbda/contacts/__init__.py,sha256=L5MM-2pv76YPGzxExdz2EErgGBATuAjYnNHlq5QOySs,503
57
- jaxsim/rbda/contacts/common.py,sha256=iywCQtesrnrwywRQv8cjyot2bG11dT_iONyF8OJztIA,5798
58
- jaxsim/rbda/contacts/relaxed_rigid.py,sha256=TR81tJ4ipcpvPnwlfkpyNDhvWizpEG542SFVu_CwHRU,19614
59
- jaxsim/rbda/contacts/rigid.py,sha256=3aDPFrIm2_QpKKRpTqJJk8qBK-W63gq7Arc8WDVAcHc,17382
60
- jaxsim/rbda/contacts/soft.py,sha256=6eFgV2hJK793RZfoY8oSqw-zC1UqFldaE0hfGHELnmU,16325
61
- jaxsim/rbda/contacts/visco_elastic.py,sha256=wATvBhLrV-7IyVLJhW7OaMg_HDAmczl_8MnYm3wuqSc,39819
57
+ jaxsim/rbda/contacts/common.py,sha256=ELASrMnkFuUZBRDPKu-HO7ymUukK3r-VZw9TjqOSZgM,11174
58
+ jaxsim/rbda/contacts/relaxed_rigid.py,sha256=6EWTkHjixSv-9pBl7aO_abOxHrneUQKF6uUyt5Vbbuc,20512
59
+ jaxsim/rbda/contacts/rigid.py,sha256=I_FrJv4rK9Lzj4i0EMaZHxOGyrPT_XuQUg6ZUzvp4y8,17943
60
+ jaxsim/rbda/contacts/soft.py,sha256=wqwnsg0abWn64pYG7je08fsBwdgyJdd2k7dAG0Sx0_k,16432
61
+ jaxsim/rbda/contacts/visco_elastic.py,sha256=S1GsYZQa8vMS1EtGmS7_Ql07h2AhJEmO67iShkGWmIU,41443
62
62
  jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
63
63
  jaxsim/terrain/terrain.py,sha256=K91HEzPqTSyNrc_j1KfAAEF_5oDeuk_-jnnZGrcMEcY,5015
64
64
  jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
65
65
  jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
66
66
  jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
67
67
  jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
68
- jaxsim-0.4.3.dev231.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
69
- jaxsim-0.4.3.dev231.dist-info/METADATA,sha256=qXdzrADStReeD5wGLpQs0reWNp1lohdSg-M_zE2hem0,17276
70
- jaxsim-0.4.3.dev231.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
71
- jaxsim-0.4.3.dev231.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
72
- jaxsim-0.4.3.dev231.dist-info/RECORD,,
68
+ jaxsim-0.4.3.dev242.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
69
+ jaxsim-0.4.3.dev242.dist-info/METADATA,sha256=RoCgOpVVtHyEo1jNXyEj6Fwq1i-Q6vodc7nqp0dT68k,17276
70
+ jaxsim-0.4.3.dev242.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
71
+ jaxsim-0.4.3.dev242.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
72
+ jaxsim-0.4.3.dev242.dist-info/RECORD,,