internal change.

PiperOrigin-RevId: 678355240
Change-Id: I271bece139de5d6e2e629d7f2fe84a4e2dd71138
This commit is contained in:
Baruch Tabanpour
2024-09-24 12:37:06 -07:00
committed by Copybara-Service
parent f0d52813b8
commit 306e52c30f
+3 -1
View File
@@ -38,6 +38,7 @@ in order to guarantee static shapes for contacts and jacobians.
""" """
import itertools import itertools
import os
from typing import Dict, Iterator, List, Tuple, Union from typing import Dict, Iterator, List, Tuple, Union
import jax import jax
@@ -393,8 +394,9 @@ def collision(m: Model, d: Data) -> Data:
# run the collision function specified by the grouping key # run the collision function specified by the grouping key
func = _COLLISION_FUNC[key.types] func = _COLLISION_FUNC[key.types]
dist, pos, frame = func(m, d, key, contact.geom)
ncon = func.ncon # pytype: disable=attribute-error ncon = func.ncon # pytype: disable=attribute-error
dist, pos, frame = func(m, d, key, contact.geom)
if ncon > 1: if ncon > 1:
# repeat contacts to match the number of collisions returned # repeat contacts to match the number of collisions returned
repeat_fn = lambda x, r=ncon: jp.repeat(x, r, axis=0) repeat_fn = lambda x, r=ncon: jp.repeat(x, r, axis=0)