Spaces:
Sleeping
Sleeping
File size: 2,994 Bytes
96da58e |
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 |
"""
Collection of useful simulation utilities
"""
from robosuite.models.base import MujocoModel
def check_contact(sim, geoms_1, geoms_2=None):
"""
Finds contact between two geom groups.
Args:
sim (MjSim): Current simulation object
geoms_1 (str or list of str or MujocoModel): an individual geom name or list of geom names or a model. If
a MujocoModel is specified, the geoms checked will be its contact_geoms
geoms_2 (str or list of str or MujocoModel or None): another individual geom name or list of geom names.
If a MujocoModel is specified, the geoms checked will be its contact_geoms. If None, will check
any collision with @geoms_1 to any other geom in the environment
Returns:
bool: True if any geom in @geoms_1 is in contact with any geom in @geoms_2.
"""
# Check if either geoms_1 or geoms_2 is a string, convert to list if so
if type(geoms_1) is str:
geoms_1 = [geoms_1]
elif isinstance(geoms_1, MujocoModel):
geoms_1 = geoms_1.contact_geoms
if type(geoms_2) is str:
geoms_2 = [geoms_2]
elif isinstance(geoms_2, MujocoModel):
geoms_2 = geoms_2.contact_geoms
for i in range(sim.data.ncon):
contact = sim.data.contact[i]
# check contact geom in geoms
c1_in_g1 = sim.model.geom_id2name(contact.geom1) in geoms_1
c2_in_g2 = sim.model.geom_id2name(contact.geom2) in geoms_2 if geoms_2 is not None else True
# check contact geom in geoms (flipped)
c2_in_g1 = sim.model.geom_id2name(contact.geom2) in geoms_1
c1_in_g2 = sim.model.geom_id2name(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):
return True
return False
def get_contacts(sim, model):
"""
Checks for any contacts with @model (as defined by @model's contact_geoms) and returns the set of
geom names currently in contact with that model (excluding the geoms that are part of the model itself).
Args:
sim (MjSim): Current simulation model
model (MujocoModel): Model to check contacts for.
Returns:
set: Unique geoms that are actively in contact with this model.
Raises:
AssertionError: [Invalid input type]
"""
# Make sure model is MujocoModel type
assert isinstance(model, MujocoModel), "Inputted model must be of type MujocoModel; got type {} instead!".format(
type(model)
)
contact_set = set()
for contact in sim.data.contact[: sim.data.ncon]:
# check contact geom in geoms; add to contact set if match is found
g1, g2 = sim.model.geom_id2name(contact.geom1), sim.model.geom_id2name(contact.geom2)
if g1 in model.contact_geoms and g2 not in model.contact_geoms:
contact_set.add(g2)
elif g2 in model.contact_geoms and g1 not in model.contact_geoms:
contact_set.add(g1)
return contact_set
|