From 306e52c30f2af5c440e2e697f7e30e6beeac874c Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 24 Sep 2024 12:37:06 -0700 Subject: [PATCH] internal change. PiperOrigin-RevId: 678355240 Change-Id: I271bece139de5d6e2e629d7f2fe84a4e2dd71138 --- mjx/mujoco/mjx/_src/collision_driver.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 8791fed7..565e586e 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -38,6 +38,7 @@ in order to guarantee static shapes for contacts and jacobians. """ import itertools +import os from typing import Dict, Iterator, List, Tuple, Union import jax @@ -393,8 +394,9 @@ def collision(m: Model, d: Data) -> Data: # run the collision function specified by the grouping key func = _COLLISION_FUNC[key.types] - dist, pos, frame = func(m, d, key, contact.geom) ncon = func.ncon # pytype: disable=attribute-error + + dist, pos, frame = func(m, d, key, contact.geom) if ncon > 1: # repeat contacts to match the number of collisions returned repeat_fn = lambda x, r=ncon: jp.repeat(x, r, axis=0)