in mae_envs/wrappers/manipulation.py [0:0]
def grab_obj(self, action):
'''
Implements object grabbing for all agents
Args:
action: Action dictionary
'''
action_pull = action['action_pull'][:, self.actual_body_slice]
sim = self.unwrapped.sim
agent_pos = sim.data.body_xpos[self.agent_body_idxs]
obj_pos = sim.data.body_xpos[self.obj_body_idxs]
obj_width = sim.model.geom_size[self.obj_geom_ids]
obj_quat = sim.data.body_xquat[self.obj_body_idxs]
assert len(obj_width) == len(obj_quat), (
"Number of object widths must be equal to number of quaternions for direct distance calculation method. " +
"This might be caused by a body that contains several geoms.")
obj_dist = dist_pt_to_cuboid(agent_pos, obj_pos, obj_width, obj_quat)
allowed_and_desired = np.logical_and(action_pull, obj_dist <= self.grab_radius)
obj_dist_masked = obj_dist.copy() # Mask the obj dists to find a valid argmin
obj_dist_masked[~allowed_and_desired] = np.inf
if self.grab_exclusive:
closest_obj = np.zeros((self.n_agents,), dtype=int)
while np.any(obj_dist_masked < np.inf):
# find agent and object of closest object distance
agent_idx, obj_idx = np.unravel_index(np.argmin(obj_dist_masked), obj_dist_masked.shape)
# set closest object for this agent
closest_obj[agent_idx] = obj_idx
# ensure exclusivity of grabbing
obj_dist_masked[:, obj_idx] = np.inf
obj_dist_masked[agent_idx, :] = np.inf
# mark same object as undesired for all other agents
allowed_and_desired[:agent_idx, obj_idx] = False
allowed_and_desired[(agent_idx + 1):, obj_idx] = False
else:
closest_obj = np.argmin(obj_dist_masked, axis=-1)
valid_grabs = np.any(allowed_and_desired, axis=-1) # (n_agent,) which agents have valid grabs
# Turn on/off agents with valid grabs
sim.model.eq_active[self.agent_eq_ids] = valid_grabs
sim.model.eq_obj2id[self.agent_eq_ids] = self.obj_body_idxs[closest_obj]
# keep track of which object is being grabbed
self.obj_grabbed = np.zeros((self.n_agents, self.n_obj), dtype=bool)
agent_with_valid_grab = np.argwhere(valid_grabs)[:, 0]
self.obj_grabbed[agent_with_valid_grab, closest_obj[agent_with_valid_grab]] = 1
# If there are new grabs, then setup the weld constraint parameters
new_grabs = np.logical_and(
valid_grabs, np.any(self.obj_grabbed != self.last_obj_grabbed, axis=-1))
for agent_idx in np.argwhere(new_grabs)[:, 0]:
agent_rot = sim.data.body_xmat[self.agent_body_idxs[agent_idx]].reshape((3, 3))
obj_rot = sim.data.body_xmat[self.obj_body_idxs[closest_obj[agent_idx]]].reshape((3, 3))
# Need to use the geom xpos rather than the qpos
obj_pos = sim.data.body_xpos[self.obj_body_idxs[closest_obj[agent_idx]]]
agent_pos = sim.data.body_xpos[self.agent_body_idxs[agent_idx]]
grab_vec = agent_pos - obj_pos
if self.grab_dist is not None:
grab_vec = self.grab_dist / (1e-3 + np.linalg.norm(grab_vec)) * grab_vec
# The distance constraint needs to be rotated into the frame of reference of the agent
sim.model.eq_data[self.agent_eq_ids[agent_idx], :3] = np.matmul(agent_rot.T, grab_vec)
# The angle constraint is the difference between the agents frame and the objects frame
sim.model.eq_data[self.agent_eq_ids[agent_idx], 3:] = mat2quat(np.matmul(agent_rot.T, obj_rot))
self.last_obj_grabbed = self.obj_grabbed