You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
71 lines
2.6 KiB
71 lines
2.6 KiB
from typing import List, Tuple
|
|
|
|
import mujoco
|
|
|
|
from decoupled_wbc.control.envs.g1.sim.sim_utilts 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
|
|
)
|