File size: 2,607 Bytes
ef6a683 c86c1a6 ef6a683 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | from typing import List, Tuple
import mujoco
from .sim_utils import get_body_geom_ids
def check_contact(
mj_model: mujoco.MjModel,
mj_data: mujoco.MjData,
bodies_1: List[str] | str,
bodies_2: List[str] | str,
return_all_contact_bodies: bool = False,
) -> Tuple[bool, List[Tuple[str, str]]] | bool:
"""
Finds contact between two body groups. Any geom in the body is considered to be in contact.
Args:
mj_model (MujocoModel): Current simulation object
mj_data (MjData): Current simulation data
bodies_1 (str or list of int): an individual body name or list of body names.
bodies_2 (str or list of int): another individual body name or list of body names.
Returns:
bool: True if any body in @bodies_1 is in contact with any body in @bodies_2.
"""
if isinstance(bodies_1, str):
bodies_1 = [bodies_1]
if isinstance(bodies_2, str):
bodies_2 = [bodies_2]
geoms_1 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_1]
geoms_1 = [g for geom_list in geoms_1 for g in geom_list]
geoms_2 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_2]
geoms_2 = [g for geom_list in geoms_2 for g in geom_list]
contact_bodies = []
for i in range(mj_data.ncon):
contact = mj_data.contact[i]
# check contact geom in geoms
c1_in_g1 = contact.geom1 in geoms_1
c2_in_g2 = contact.geom2 in geoms_2 if geoms_2 is not None else True
# check contact geom in geoms (flipped)
c2_in_g1 = contact.geom2 in geoms_1
c1_in_g2 = contact.geom1 in geoms_2 if geoms_2 is not None else True
if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1):
contact_bodies.append(
(
mj_model.body(mj_model.geom(contact.geom1).bodyid).name,
mj_model.body(mj_model.geom(contact.geom2).bodyid).name,
)
)
if not return_all_contact_bodies:
break
if return_all_contact_bodies:
return len(contact_bodies) > 0, set(contact_bodies)
else:
return len(contact_bodies) > 0
def check_height(
mj_model: mujoco.MjModel,
mj_data: mujoco.MjData,
geom_name: str,
lower_bound: float = -float("inf"),
upper_bound: float = float("inf"),
):
"""
Checks if the height of a geom is greater than a given height.
"""
geom_id = mj_model.geom(geom_name).id
return (
mj_data.geom_xpos[geom_id][2] < upper_bound and mj_data.geom_xpos[geom_id][2] > lower_bound
)
|